From 106b2934dc17ca7b0a4b16192cfbc5496e8bbfad Mon Sep 17 00:00:00 2001 From: Jithun Nair Date: Tue, 17 Aug 2021 16:25:52 +0000 Subject: [PATCH] Make torch version check numeric --- setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 60fd4ac3e0b..70570b0935c 100644 --- 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