Update composer/trainer/_patch_pytorch.py

Co-authored-by: Mihir Patel <[email protected]>
mosaicml · Sep 10, 2024 · f5e0398 · f5e0398
1 parent 2664be4
commit f5e0398
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/composer/trainer/_patch_pytorch.py b/composer/trainer/_patch_pytorch.py
@@ -943,7 +943,9 @@ def unshard_with_sync(self):
         self._use_unsharded_flat_param(padded_unsharded_flat_param)
 
 
-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.4.1'):
     # 2.4.0 only patch
     # PyTorch issue: https://github.com/pytorch/pytorch/issues/133923
     from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE