Make torch version check numeric (#4285)

Co-authored-by: Vasilis Vryniotis <[email protected]>
pytorch · Aug 17, 2021 · f3aff2f · f3aff2f
1 parent af16236
commit f3aff2f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -147,7 +147,9 @@ def get_extensions():
     )
 
     is_rocm_pytorch = False
-    if torch.__version__ >= '1.5':
+    TORCH_MAJOR = int(torch.__version__.split('.')[0])
+    TORCH_MINOR = int(torch.__version__.split('.')[1])
+    if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5):
         from torch.utils.cpp_extension import ROCM_HOME
         is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False