fix training in cascade_rcnn (PaddlePaddle#2386)

heavengate · Mar 19, 2021 · a2f57be · a2f57be
1 parent 588d78b
commit a2f57be
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/dygraph/ppdet/modeling/heads/cascade_head.py b/dygraph/ppdet/modeling/heads/cascade_head.py
@@ -227,6 +227,8 @@ def _get_rois_from_boxes(self, boxes, im_shape):
             clip_box = clip_bbox(boxes_per_image, im_shape[i])
             if self.training:
                 keep = nonempty_bbox(clip_box)
+                if keep.shape[0] == 0:
+                    keep = paddle.zeros([1], dtype='int32')
                 clip_box = paddle.gather(clip_box, keep)
             rois.append(clip_box)
         rois_num = paddle.concat([paddle.shape(r)[0] for r in rois])