diff --git a/mmcv/ops/bbox.py b/mmcv/ops/bbox.py index 4ba93d6b2..4583ba7d5 100644 --- a/mmcv/ops/bbox.py +++ b/mmcv/ops/bbox.py @@ -116,6 +116,10 @@ def bbox_overlaps(bboxes1: torch.Tensor, if rows * cols == 0: return ious + if bboxes1.device.type == 'cpu' and torch.__version__ == 'parrots': + return _bbox_overlaps_cpu( + bboxes1, bboxes2, mode=mode, aligned=aligned, offset=offset) + ext_module.bbox_overlaps( bboxes1, bboxes2, ious, mode=mode_flag, aligned=aligned, offset=offset) diff --git a/setup.py b/setup.py index 176db6b87..c12a85358 100644 --- a/setup.py +++ b/setup.py @@ -212,6 +212,7 @@ def get_extensions(): glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') +\ glob.glob('./mmcv/ops/csrc/parrots/*.cpp') op_files.remove('./mmcv/ops/csrc/pytorch/cuda/iou3d_cuda.cu') + op_files.remove('./mmcv/ops/csrc/pytorch/cpu/bbox_overlaps_cpu.cpp') include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common')) include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda')) cuda_args = os.getenv('MMCV_CUDA_ARGS')