Pin checkpoint state dict flattening patch

mosaicml · Nov 4, 2024 · da96c16 · da96c16
1 parent 969d3b1
commit da96c16
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -595,7 +595,7 @@ 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(