Begin implementing a training edge case to pickup where the training …

…left off if interrupeted.
jshuadvd · Jul 7, 2024 · 32e537f · 32e537f
1 parent c70b0bc
commit 32e537f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -282,7 +282,9 @@ def train(
             best_val_loss = avg_val_loss
             patience = 0
             # Save best model
-            accelerator.save_state("best_model.pt")
+            accelerator.save_state(
+                {"epoch": epoch, "best_val_loss": best_val_loss}, "best_model.pt"
+            )
         else:
             patience += 1
             if patience >= max_patience: