diff --git a/setup.py b/setup.py index 12f3f46b86..a0b6a429cc 100644 --- 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 index 3d1486eb01..752d877663 100644 --- 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(