bugfixed: remove batch_index after proposal_layer

0d524bc1 · Shaoshuai Shi · 202d767d · 0d524bc1
Commit 0d524bc1 authored Jul 23, 2020 by Shaoshuai Shi
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

pcdet/models/roi_heads/roi_head_template.py pcdet/models/roi_heads/roi_head_template.py +1 -0

No files found.
--- a/pcdet/models/roi_heads/roi_head_template.py
+++ b/pcdet/models/roi_heads/roi_head_template.py
@@ -92,6 +92,7 @@ class RoIHeadTemplate(nn.Module):
        batch_dict['roi_scores'] = roi_scores
        batch_dict['roi_labels'] = roi_labels + 1
        batch_dict['has_class_labels'] = True if batch_cls_preds.shape[-1] > 1 else False
+        batch_dict.pop('batch_index', None)
        return batch_dict

    def assign_targets(self, batch_dict):