Add DCP Saver module to the callbacks target (#775)

Summary: Pull Request resolved: #775 Reviewed By: JKSenthil Differential Revision: D55618741 fbshipit-source-id: 46e8192ca8777450d111dd5570d13d2ccb2e7d2e
pytorch · Apr 3, 2024 · ec33f1b · ec33f1b
1 parent c01862b
commit ec33f1b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchtnt/framework/callbacks/__init__.py b/torchtnt/framework/callbacks/__init__.py
@@ -7,6 +7,7 @@
 # pyre-strict
 
 from .base_csv_writer import BaseCSVWriter
+from .dcp_saver import DistributedCheckpointSaver
 from .early_stopping import EarlyStopping
 from .empty_cuda_cache import EmptyCudaCache
 from .garbage_collector import GarbageCollector
@@ -44,4 +45,5 @@
     "TorchSnapshotSaver",
     "TQDMProgressBar",
     "TrainProgressMonitor",
+    "DistributedCheckpointSaver",
 ]