Make torch version check numeric (#4285)

jithunnair-amd · datumbox · web-flow · commit f3aff2fa60c1 · 2021-08-17T18:01:38.000+01:00
Co-authored-by: Vasilis Vryniotis &lt;datumbox@users.noreply.github.com&gt;
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
 

Original file line number	Diff line number	Diff line change
`@@ -147,7 +147,9 @@ def get_extensions():`
`147`	`147`	`)`
`148`	`148`
`149`	`149`	`is_rocm_pytorch = False`
`150`		`- if torch.__version__ >= '1.5':`
	`150`	`+ TORCH_MAJOR = int(torch.__version__.split('.')[0])`
	`151`	`+ TORCH_MINOR = int(torch.__version__.split('.')[1])`
	`152`	`+ if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5):`
`151`	`153`	`from torch.utils.cpp_extension import ROCM_HOME`
`152`	`154`	`is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False`
`153`	`155`