From ef075060c0d2be4c4b28ae0c2a8e87d3842c6519 Mon Sep 17 00:00:00 2001 From: ABD-01 Date: Mon, 23 Aug 2021 09:35:40 +0000 Subject: [PATCH] Updated setup.py to use TorchVersion object for version comparision --- setup.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 70570b0935c..ca3db84eb19 100644 --- a/setup.py +++ b/setup.py @@ -147,9 +147,7 @@ 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