Merge pull request #5 from fsherry/main

maximilianherde · web-flow · commit 2cc8c85f6184 · 2024-06-20T13:54:51.000+02:00
Explicitly save model after training (to deal with bug in transformers 4.29.2)
diff --git a/scOT/train.py b/scOT/train.py
@@ -407,7 +407,8 @@ def get_statistics(errors):
     )
 
     trainer.train(resume_from_checkpoint=params.resume_training)
-
+    trainer.save_model(train_config.output_dir)
+    
     if (RANK == 0 or RANK == -1) and params.push_to_hf_hub is not None:
         model.push_to_hub(params.push_to_hf_hub)
 

Original file line number	Diff line number	Diff line change
`@@ -407,7 +407,8 @@ def get_statistics(errors):`
`407`	`407`	`)`
`408`	`408`
`409`	`409`	`trainer.train(resume_from_checkpoint=params.resume_training)`
`410`		`-`
	`410`	`+ trainer.save_model(train_config.output_dir)`
	`411`	`+`
`411`	`412`	`if (RANK == 0 or RANK == -1) and params.push_to_hf_hub is not None:`
`412`	`413`	`model.push_to_hub(params.push_to_hf_hub)`
`413`	`414`