Improve coverage

NicolasHug · NicolasHug · commit 4b7e94229498 · 2021-03-08T10:03:54.000Z
diff --git a/test/test_ops.py b/test/test_ops.py
@@ -479,6 +479,10 @@ def test_batched_nms_implementations(self):
         err_msg = "The vanilla and the trick implementation yield different nms outputs."
         self.assertTrue(torch.allclose(keep_vanilla, keep_trick), err_msg)
 
+        # Also make sure an empty tensor is returned if boxes is empty
+        empty = torch.empty((0,), dtype=torch.int64)
+        self.assertTrue(torch.allclose(empty, ops.batched_nms(empty, None, None, None)))
+
 
 class DeformConvTester(OpTester, unittest.TestCase):
     def expected_fn(self, x, weight, offset, mask, bias, stride=1, padding=0, dilation=1):
diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py
@@ -98,9 +98,6 @@ def _batched_nms_vanilla(
     iou_threshold: float,
 ) -> Tensor:
     # Based on Detectron2 implementation, just manually call nms() on each class independently
-    if boxes.numel() == 0:
-        return torch.empty((0,), dtype=torch.int64, device=boxes.device)
-
     keep_mask = torch.zeros_like(scores, dtype=torch.bool)
     for class_id in torch.unique(idxs):
         curr_indices = torch.where(idxs == class_id)[0]
