Update composer/callbacks/oom_observer.py

Co-authored-by: Mihir Patel <[email protected]>
mosaicml · Feb 2, 2024 · 5a74d34 · 5a74d34
1 parent 43118ca
commit 5a74d34
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/composer/callbacks/oom_observer.py b/composer/callbacks/oom_observer.py
@@ -86,7 +86,7 @@ def __init__(
         else:
             self.remote_path_in_bucket = None
 
-        if version.parse(torch.__version__) > version.parse('2.1.0.dev'):  # type: ignore
+        if version.parse(torch.__version__.split('.dev')[0]) >= version.parse('2.1.0'):  # type: ignore
             # OOMObserver  is only supported in torch v2.1.0-rc1 or higher
             self._enabled = True
         else: