temporarily disable mps ops for torch2.1.0 (#2958)

open-mmlab · Dec 26, 2023 · 1bb3dea · 1bb3dea
1 parent 62d7092
commit 1bb3dea
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -424,9 +424,11 @@ def get_mluops_version(file_path):
             extra_compile_args['cxx'] += ['-ObjC++']
             # src
             op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
-                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
-                glob.glob('./mmcv/ops/csrc/common/mps/*.mm') + \
-                glob.glob('./mmcv/ops/csrc/pytorch/mps/*.mm')
+                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp')
+            # TODO: support mps ops on torch>=2.1.0
+            if parse_version(torch.__version__) < parse_version('2.1.0'):
+                op_files += glob.glob('./mmcv/ops/csrc/common/mps/*.mm') + \
+                    glob.glob('./mmcv/ops/csrc/pytorch/mps/*.mm')
             extension = CppExtension
             include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
             include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/mps'))

diff --git a/tests/test_ops/test_bbox.py b/tests/test_ops/test_bbox.py
@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 import torch
+from mmengine.utils import digit_version
 
 from mmcv.utils import (IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE, IS_MPS_AVAILABLE,
                         IS_NPU_AVAILABLE)
@@ -56,7 +57,9 @@ def _test_bbox_overlaps(self, device='cpu', dtype=torch.float):
         pytest.param(
             'mps',
             marks=pytest.mark.skipif(
-                not IS_MPS_AVAILABLE, reason='requires MPS support')),
+                not IS_MPS_AVAILABLE
+                or digit_version(torch.__version__) >= digit_version('2.1.0'),
+                reason='requires MPS support')),
         pytest.param(
             'npu',
             marks=pytest.mark.skipif(