Update setup.py

Cambricon · Jul 9, 2024 · cfbdfea · cfbdfea
1 parent 563cc2e
commit cfbdfea
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
     if torch.__version__ == 'parrots':
         from parrots.utils.build_extension import BuildExtension
         EXT_TYPE = 'parrots'
-    elif (hasattr(torch, 'mlu') and torch.mlu.is_available()) or \
+    elif (hasattr(torch, 'mlu')) or \
             os.getenv('FORCE_MLU', '0') == '1':
         from torch_mlu.utils.cpp_extension import BuildExtension
         EXT_TYPE = 'pytorch'
@@ -50,7 +50,7 @@ def get_version():
     return locals()['__version__']
 
 def get_mlu_version():
-    if not (hasattr(torch, 'mlu') and torch.mlu.is_available()):
+    if not (hasattr(torch, 'mlu')):
         return "", ""
     version_file = './build.property'
     import json
@@ -338,8 +338,7 @@ def get_extensions():
             include_dirs.append(os.path.abspath('./mmcv/ops/csrc/pytorch'))
             include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
             include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))
-        elif (hasattr(torch, 'mlu') and
-                torch.mlu.is_available()) or \
+        elif (hasattr(torch, 'mlu')) or \
                 os.getenv('FORCE_MLU', '0') == '1':
             from torch_mlu.utils.cpp_extension import MLUExtension