Missing kernel init on Impala convolution layers wanted?

[theovincent]

Hi,

Thanks again for the amazing code base!

The second and third definitions of the convolution layers are missing the initializer type, while the initializer type is defined earlier. Is that wanted?

dopamine/dopamine/jax/networks.py

Lines 87 to 89 in bec5f4e

	conv_out = nn.Conv(
	features=self.num_ch, kernel_size=(3, 3), strides=1, padding='SAME'
	)(conv_out)

Looking at the original code of BBF, the kernel init was specified for the same convolution layers:
https://github.com/google-research/google-research/blob/89dbaf2657e70797f2dea7a6acac5a808da8ff4d/bigger_better_faster/bbf/spr_networks.py#L557

Best,
Théo Vincent