test

run_mlm_flax_stream.py

@@ -660,6 +660,10 @@ if __name__ == "__main__":
 
                     model_pt.push_to_hub(training_args.hub_model_id,commit_message=f"Weights for torch of step {step+1}")
                     model_flax.push_to_hub(training_args.hub_model_id,commit_message=f"Weights for flax of step {step+1}")
+                    
+                    #Delete the models to free memory
+                    del(model_pt)
+                    del(model_flax)
 
                     print(f"Saving weights and logs of step {step+1}. \nThe result is saved to {training_args.output_dir} by worker {jax.process_index()}.")