Enable DCP checkpoints with on_eval_epoch_end hook (#851)

Summary: Pull Request resolved: #851

Reviewed By: zedsdead01, JKSenthil

Differential Revision: D58967392

fbshipit-source-id: 0ecdd0bec0804167127daa2530bf77b63bf7ca6e

Author: diego-urgell

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"