Fix getting batch size

open-mmlab · Apr 30, 2021 · b0d5be6 · b0d5be6
1 parent 3f7632b
commit b0d5be6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmdet/models/roi_heads/point_rend_roi_head.py b/mmdet/models/roi_heads/point_rend_roi_head.py
@@ -69,7 +69,7 @@ def _get_fine_grained_point_feats(self, x, rois, rel_roi_points,
  feats = x[idx]
  spatial_scale = 1. / float(
  self.mask_roi_extractor.featmap_strides[idx])
- if torch.onnx.is_in_onnx_export() and num_imgs == 1:
+ if torch.onnx.is_in_onnx_export() and rois.shape[0] == 1:
  feat = feats[0].unsqueeze(0)
  rel_img_points = rel_roi_point_to_rel_img_point(
  rois, rel_roi_points, feat, spatial_scale).unsqueeze(0)