Fix torch version check (#3161)

* fix * comment
mosaicml · Apr 2, 2024 · 4d88824 · 4d88824
1 parent 31f1159
commit 4d88824
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/composer/trainer/mosaic_fsdp.py b/composer/trainer/mosaic_fsdp.py
@@ -61,8 +61,8 @@ def patch_pytorch():
         from torch.distributed.fsdp import _runtime_utils
         _runtime_utils._validate_and_get_hybrid_shard_state = lambda *args, **kwargs: None
 
-    elif version.parse(torch.__version__) < version.parse('2.3.0'):
-        # Monkey patch for torch < 2.2.2 ie torch == 2.2.1
+    elif version.parse(torch.__version__) < version.parse('2.2.9'):
+        # Monkey patch for torch < 2.3.0 ie torch == 2.2.1/2.2.2 currently
         pass
 
     elif version.parse(torch.__version__) < version.parse('2.3.1'):