Skip to content

GitLab

Unverified Commit 221dd5f7 authored by Muhammed Abdullah's avatar Muhammed Abdullah Committed by GitHub
Browse files

Updated setup.py to use TorchVersion object for version comparision (#4307)

parent a9940fe4
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
setup.py
View file @ 221dd5f7
...@@ -139,9 +139,8 @@ def get_extensions(): ...@@ -139,9 +139,8 @@ def get_extensions():
) )
is_rocm_pytorch = False is_rocm_pytorch = False
TORCH_MAJOR = int(torch.__version__.split(".")[0])
TORCH_MINOR = int(torch.__version__.split(".")[1]) if torch.__version__ >= "1.5":
if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5):
from torch.utils.cpp_extension import ROCM_HOME 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 is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment