Changed batch size

run.sh

@@ -6,10 +6,16 @@
     --dataset_name="large_spanish_corpus" \
     --dataset_config_name="combined" \
     --max_seq_length="128" \
-    --per_device_train_batch_size="4" \
-    --per_device_eval_batch_size="4" \
+    --pad_to_max_length  \
+    --per_device_train_batch_size="128" \
+    --per_device_eval_batch_size="128" \
     --learning_rate="3e-4" \
+    --save_strategy="steps" \
+    --save_steps="10000" \
+    --save_total_limit="5" \
     --warmup_steps="1000" \
     --overwrite_output_dir \
     --num_train_epochs="8" \
-    --push_to_hub
+    --dtype="bfloat16" \
+    --push_to_hub_model_id="flax-community/bertin-roberta-large-spanish" \
+    --push_to_hub 2>&1 | tee run.log

run_mlm_flax.py

@@ -315,10 +315,10 @@ if __name__ == "__main__":
 
     # Log on each process the small summary:
     logger = logging.getLogger(__name__)
-    logger.warning(
-        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
-        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
-    )
+    #logger.warning(
+    #    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+    #    + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    #)
 
     # Set the verbosity to info of the Transformers logger (on main process only):
     logger.info(f"Training/evaluation parameters {training_args}")