fix rtmdet-ins head (#1884)

2023-03-20 12:03:12 +08:00 · 2023-03-20 12:03:12 +08:00 · 754c08012d
parent dce6e6401e
commit 754c08012d
1 changed files with 3 additions and 0 deletions
--- a/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py
+++ b/mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_ins_head.py
@ -170,6 +170,9 @@ def _nms_with_mask_static(self,
    # pad
    dets = torch.cat((dets, dets.new_zeros((1, 1, 5))), 1)
    labels = torch.cat((labels, labels.new_zeros((1, 1))), 1)
+    kernels = torch.cat((kernels, kernels.new_zeros(1, 1, kernels.shape[2])),
+                        1)
+    priors = torch.cat((priors, priors.new_zeros(1, 4)), 0)

    # topk or sort
    is_use_topk = keep_top_k > 0 and \