Machine-Learning
/
aiuNN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixed saving by adding a save if there is no improvement
Gitea Actions For AIIA / Explore-Gitea-Actions (push) Successful in 10m52s Details

This commit is contained in:
Falko Victor Habel 2025-06-30 09:54:07 +02:00
parent 1cf8b9e09b
commit 358a7c34d0
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/aiunn/finetune/trainer.py
View File

@ -309,7 +309,9 @@ class aiuNNTrainer:
# Cleanup # Cleanup
gc.collect() gc.collect()
torch.cuda.empty_cache() 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 return self.best_loss
def load_checkpoint(self, specific_checkpoint=None): def load_checkpoint(self, specific_checkpoint=None):