Merge pull request #672 from runame/fix-rng

Use separate `rng` for dropout in FastMRI `model_init`

Author: priyakasimbeg

algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py

@@ -33,8 +33,10 @@ def init_model_fn(
         use_tanh=self.use_tanh,
         use_layer_norm=self.use_layer_norm,
         dropout_rate=dropout_rate)
-
-    variables = jax.jit(self._model.init)({'params': rng}, fake_batch)
+    params_rng, dropout_rng = jax.random.split(rng)
+    variables = jax.jit(
+        self._model.init)({'params': params_rng, 'dropout': dropout_rng},
+                          fake_batch)
     params = variables['params']
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)