added saving latest model when keyboard interrupt during training laat

pokarats · Jan 24, 2023 · bc87e1c · bc87e1c
1 parent 1915540
commit bc87e1c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/models/train_eval_laat.py b/src/models/train_eval_laat.py
@@ -131,6 +131,10 @@ def train(
     except KeyboardInterrupt:
         print('*' * 20)
         print('Exiting from training early')
+        if not (MODEL_FOLDER / f"best_{model_save_fname}.pt").exists():
+            logger.info(f"saving best model so far from current epoch...")
+            torch.save(model.state_dict(), f"{MODEL_FOLDER / f'best_{model_save_fname}.pt'}")
+        return evals
 
     return evals