diff --git a/composer/trainer/_patch_pytorch.py b/composer/trainer/_patch_pytorch.py
index 4cd81a09eb..fcca94d73a 100644
--- 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