pere
/

roberta-base-exp-32B

Model card Files Files and versions

pere commited on Nov 17, 2022

Commit

dfcb939

·

1 Parent(s): 7d84292

Update run_mlm_flax_stream.py

Files changed (1) hide show

run_mlm_flax_stream.py +2 -2

run_mlm_flax_stream.py CHANGED Viewed

@@ -449,8 +449,8 @@ if __name__ == "__main__":
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count() * jax.process_count()
-    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count() * jax.process_count()
     # define number steps per stream epoch
     num_train_steps = data_args.num_train_steps

     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
     # define number steps per stream epoch
     num_train_steps = data_args.num_train_steps