Skip to content

GitLab

Unverified Commit f3aff2fa authored by Jithun Nair's avatar Jithun Nair Committed by GitHub
Browse files

Make torch version check numeric (#4285) 


Co-authored-by: default avatarVasilis Vryniotis <datumbox@users.noreply.github.com>
parent af16236d
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
setup.py
View file @ f3aff2fa
......@@ -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
......
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