fixed saving by adding a save if there is no improvement

2025-06-30 09:54:07 +02:00 · 2025-06-30 09:54:07 +02:00 · 358a7c34d0
parent 1cf8b9e09b
commit 358a7c34d0
1 changed files with 3 additions and 1 deletions
--- a/src/aiunn/finetune/trainer.py
+++ b/src/aiunn/finetune/trainer.py
@ -309,7 +309,9 @@ class aiuNNTrainer:
            # Cleanup
            gc.collect()
            torch.cuda.empty_cache()
-        
+        if not is_improved:
+            best_model_path = os.path.join(output_path, "final_model")
+            self.model.save_pretrained(best_model_path)
        return self.best_loss

    def load_checkpoint(self, specific_checkpoint=None):