Pin checkpoint state dict flattening patch (#3700)

mosaicml · Nov 4, 2024 · 03c992b · 03c992b
1 parent 969d3b1
commit 03c992b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -595,7 +595,10 @@ def dist_cp_load(
     storage_reader: StorageReader,
     load_planner: Optional[LoadPlanner] = None,
 ):
-    if version.parse(torch.__version__) >= version.parse('2.4.0'):
+    if (
+        version.parse(torch.__version__) >= version.parse('2.4.0') and
+        version.parse(torch.__version__) < version.parse('2.5.0')
+    ):
         from torch.distributed.checkpoint.utils import CheckpointException
         try:
             dist_cp.load(