Enable DCP checkpoints with on_eval_epoch_end hook (pytorch#851)

Summary: Pull Request resolved: pytorch#851 Reviewed By: zedsdead01, JKSenthil Differential Revision: D58967392 fbshipit-source-id: 0ecdd0bec0804167127daa2530bf77b63bf7ca6e
diego-urgell · Jun 25, 2024 · 1d118c4 · 1d118c4
1 parent 8afe26d
commit 1d118c4
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/torchtnt/framework/callbacks/dcp_saver.py b/torchtnt/framework/callbacks/dcp_saver.py
@@ -136,7 +136,12 @@ def _checkpoint_impl(
         planner: Optional[SavePlanner] = None,
         storage_writer: Optional[StorageWriter] = None,
     ) -> bool:
-        if hook not in ["on_train_step_end", "on_train_epoch_end", "on_train_end"]:
+        if hook not in [
+            "on_train_step_end",
+            "on_train_epoch_end",
+            "on_train_end",
+            "on_eval_epoch_end",
+        ]:
             raise RuntimeError(f"Unexpected hook encountered '{hook}'")
 
         intra_epoch = hook == "on_train_step_end"