pere
/

norwegian-bert-base-newcode

Model card Files Files and versions Metrics Training metrics Community

pere commited on 6 days ago

Commit

b69bc4d

•

1 Parent(s): 711f1cb

Update run_mlm_flax.py

Files changed (1) hide show

run_mlm_flax.py +2 -1

run_mlm_flax.py CHANGED Viewed

@@ -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,))
     # Define eval fn
     def eval_step(params, batch):

         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")
     # Define eval fn
     def eval_step(params, batch):