diff --git a/mmdeploy/mmcv/ops/nms.py b/mmdeploy/mmcv/ops/nms.py index b384d75d47..952a17b6ef 100644 --- a/mmdeploy/mmcv/ops/nms.py +++ b/mmdeploy/mmcv/ops/nms.py @@ -349,10 +349,6 @@ def _multiclass_nms_single(boxes: Tensor, dets = torch.cat([boxes, scores], dim=2) labels = cls_inds.unsqueeze(0) - # pad - dets = torch.cat((dets, dets.new_zeros((1, 1, 5))), 1) - labels = torch.cat((labels, labels.new_zeros((1, 1))), 1) - # topk or sort is_use_topk = keep_top_k > 0 and \ (torch.onnx.is_in_onnx_export() or keep_top_k < dets.shape[1]) @@ -369,7 +365,7 @@ def _multiclass_nms_single(boxes: Tensor, if pre_top_k > 0: bbox_index = pre_topk_inds[None, box_inds] if keep_top_k > 0: - bbox_index = bbox_index[:, topk_inds[:-1]] + bbox_index = bbox_index[:, topk_inds] return dets, labels, bbox_index else: return dets, labels