Update run_mlm_flax.py

run_mlm_flax.py

@@ -786,7 +786,8 @@ def main():
         return new_state, metrics, new_dropout_rng
 
     # Create parallel version of the train step
-    p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+    # p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+    p_train_step = jax.pmap(train_step, "batch")
 
     # Define eval fn
     def eval_step(params, batch):