From cfd3a3a70738d3336dac4a348571871a24958c06 Mon Sep 17 00:00:00 2001
From: VVsssssk <88368822+VVsssssk@users.noreply.github.com>
Date: Fri, 30 Sep 2022 16:36:35 +0800
Subject: [PATCH] [Fix]Fix pointrcnn in dev branch (#1874)

* fix parta2 bug

* fix

* fix
---
 mmdet3d/models/dense_heads/point_rpn_head.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mmdet3d/models/dense_heads/point_rpn_head.py b/mmdet3d/models/dense_heads/point_rpn_head.py
index 546cf1665..8704acb92 100644
--- a/mmdet3d/models/dense_heads/point_rpn_head.py
+++ b/mmdet3d/models/dense_heads/point_rpn_head.py
@@ -271,9 +271,13 @@ def get_bboxes(self,
         for b in range(batch_size):
             bbox3d = self.bbox_coder.decode(bbox_preds[b], points[b, ..., :3],
                                             object_class[b])
+            mask = ~bbox3d.sum(dim=1).isinf()
             bbox_selected, score_selected, labels, cls_preds_selected = \
-                self.class_agnostic_nms(obj_scores[b], sem_scores[b], bbox3d,
-                                        points[b, ..., :3], input_metas[b])
+                self.class_agnostic_nms(obj_scores[b][mask],
+                                        sem_scores[b][mask, :],
+                                        bbox3d[mask, :],
+                                        points[b, ..., :3][mask, :],
+                                        input_metas[b])
             bbox = input_metas[b]['box_type_3d'](
                 bbox_selected.clone(),
                 box_dim=bbox_selected.shape[-1],