From 2c4475e3aad389bd5a8ebd48eb383b27e5eb5e0e Mon Sep 17 00:00:00 2001 From: Rithwik Ediga Lakhamsani Date: Wed, 19 Feb 2025 11:54:14 -0800 Subject: [PATCH] removed upper bound for dist_cp.load_state_dict --- composer/utils/checkpoint.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py index 4dd1910a4a..8e910c81e3 100644 --- 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(