Add dropout_rng in fastmri model_init

Author: runame

algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py

@@ -33,8 +33,9 @@ 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)