Merge pull request #13 from B-lanc/bug-train_state

moved epoch state increment to before save
f90 · May 8, 2022 · 86c113d · 86c113d
2 parents 443fe30 + 0e13994
commit 86c113d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -134,11 +134,11 @@ def main(args):
             state["best_loss"] = val_loss
             state["best_checkpoint"] = checkpoint_path
 
+        state["epochs"] += 1
         # CHECKPOINT
         print("Saving model...")
         model_utils.save_model(model, optimizer, state, checkpoint_path)
 
-        state["epochs"] += 1
 
     #### TESTING ####
     # Test loss