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

…left off if interrupeted.
jshuadvd · Jul 7, 2024 · 2fe7e9e · 2fe7e9e
1 parent b7d2497
commit 2fe7e9e
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -267,7 +267,15 @@ def train(
         )
 
         # Save checkpoint
-        accelerator.save_state(f"checkpoint_epoch_{epoch}.pt")
+        accelerator.save_state(
+            {"epoch": epoch, "best_val_loss": best_val_loss},
+            f"checkpoint_epoch_{epoch}.pt",
+        )
+
+        # Save latest checkpoint
+        accelerator.save_state(
+            {"epoch": epoch, "best_val_loss": best_val_loss}, "checkpoint_latest.pt"
+        )
 
         # Early stopping
         if avg_val_loss < best_val_loss: