add nms_rotated adpated torch2.1

momo609 · Sep 14, 2023 · 56111ee · 56111ee
1 parent b361a81
commit 56111ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmcv/ops/csrc/pytorch/nms_rotated.cpp b/mmcv/ops/csrc/pytorch/nms_rotated.cpp
@@ -36,7 +36,7 @@ Tensor nms_rotated(const Tensor dets, const Tensor scores, const Tensor order,
 #else
     AT_ERROR("Not compiled with GPU support");
 #endif
-  } else if (dets.device().type() == at::kXLA) {
+  } else if (dets.device().type() == at::kXLA || dets.device().type() == at::kPrivateUse1) {
 #ifdef MMCV_WITH_NPU
     return nms_rotated_npu(dets, scores, labels, iou_threshold);
 #else