removed upper bound for dist_cp.load_state_dict

mosaicml · Feb 19, 2025 · 2c4475e · 2c4475e
1 parent 5e1b16e
commit 2c4475e
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -557,10 +557,7 @@ def dist_cp_load(
     storage_reader: StorageReader,
     load_planner: Optional[LoadPlanner] = None,
 ):
-    if (
-        version.parse(torch.__version__) >= version.parse('2.4.0') and
-        version.parse(torch.__version__) < version.parse('2.5.0') # TODO(Check): It seems that we shouldn't have a higher bound here.
-    ):
+    if version.parse(torch.__version__) >= version.parse('2.4.0'):
         from torch.distributed.checkpoint.utils import CheckpointException
         try:
             dist_cp.load(