new attempt

run.sh

@@ -1,7 +1,6 @@
 
 python run_mlm_flax_stream.py \
     --output_dir="./" \
-    --hub_model_id="NbAiLab/nb-roberta-base-scandi" \
     --model_name_or_path="xlm-roberta-base" \
     --config_name="./config.json" \
     --tokenizer_name="./" \
@@ -13,6 +12,7 @@ python run_mlm_flax_stream.py \
     --learning_rate="4e-4" \
     --warmup_steps="1000" \
     --overwrite_output_dir \
+    --use_auth_token \
     --num_train_steps="10000" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
@@ -20,4 +20,5 @@ python run_mlm_flax_stream.py \
     --save_steps="50" \
     --eval_steps="50" \
     --dtype="bfloat16" \
+    --push_to_hub_model_id="NbAiLab/nb-roberta-base-scandi" \
     --push_to_hub

run_mlm_flax_stream.py

@@ -655,13 +655,18 @@ if __name__ == "__main__":
                     )
                     print(f"Saving weights and logs of step {step+1}. \nThe result is saved to {training_args.output_folder} by worker {jax.process_index()}.")
                 except:
-
+                    breakpoint()
                     model.save_pretrained(
                         training_args.output_dir,
                         params=params
                     )
                     print("Problems pushing this to the hub. The bug should be fixed.")
-
+            else:
+                model.save_pretrained(
+                    training_args.output_dir,
+                    params=params
+                )
+            
         # update tqdm bar
         steps.update(1)