test

run.sh

@@ -1,13 +1,13 @@
 python run_mlm_flax_stream.py \
-    --output_dir="../roberta-debug-32" \
+    --output_dir="../roberta-debug-pod-32" \
     --model_name_or_path="xlm-roberta-base" \
     --config_name="./" \
     --tokenizer_name="./" \
     --dataset_name="NbAiLab/scandinavian" \
     --max_seq_length="512" \
     --weight_decay="0.01" \
-    --per_device_train_batch_size="12" \
-    --per_device_eval_batch_size="12" \
+    --per_device_train_batch_size="62" \
+    --per_device_eval_batch_size="62" \
     --learning_rate="1e-4" \
     --warmup_steps="10000" \
     --overwrite_output_dir \

run_mlm_flax_stream.py

@@ -451,7 +451,7 @@ if __name__ == "__main__":
     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()
-    
+    breakpoint() 
     print("***************************")
     print(f"Train Batch Size: {train_batch_size}")
     print("***************************")