Fixed 0-elements bug in forward (facebookresearch#244)

bjuncek · Sep 23, 2020 · 7cd7895 · 7cd7895
1 parent 8889f08
commit 7cd7895
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/segmentation.py b/models/segmentation.py
@@ -284,7 +284,7 @@ def to_tuple(tup):
             cur_scores = cur_scores[keep]
             cur_classes = cur_classes[keep]
             cur_masks = cur_masks[keep]
-            cur_masks = interpolate(cur_masks[None], to_tuple(size), mode="bilinear").squeeze(0)
+            cur_masks = interpolate(cur_masks[:, None], to_tuple(size), mode="bilinear").squeeze(1)
             cur_boxes = box_ops.box_cxcywh_to_xyxy(cur_boxes[keep])
 
             h, w = cur_masks.shape[-2:]