[fbsync] Updated setup.py to use TorchVersion object for version comp…

…arision (#4307) Reviewed By: fmassa Differential Revision: D31649964 fbshipit-source-id: 670d26caf53ffc9127ed22afc74cec8ec5641595
pytorch · Oct 14, 2021 · fd58d13 · fd58d13
1 parent 5a270a5
commit fd58d13
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -139,9 +139,8 @@ def get_extensions():
     )
 
     is_rocm_pytorch = False
-    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):
+
+    if torch.__version__ >= "1.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