diff --git a/test/test_ops.py b/test/test_ops.py index 914dd93b988..4945c7b9714 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -374,10 +374,14 @@ def _create_tensors_with_iou(self, N, iou_thresh): # let b0 be [x0, y0, x1, y1], and b1 be [x0, y0, x1 + d, y1], # then, in order to satisfy ops.iou(b0, b1) == iou_thresh, # we need to have d = (x1 - x0) * (1 - iou_thresh) / iou_thresh + # Adjust the threshold upward a bit with the intent of creating + # at least one box that exceeds (barely) the threshold and so + # should be suppressed. boxes = torch.rand(N, 4) * 100 boxes[:, 2:] += boxes[:, :2] boxes[-1, :] = boxes[0, :] x0, y0, x1, y1 = boxes[-1].tolist() + iou_thresh += 1e-5 boxes[-1, 2] += (x1 - x0) * (1 - iou_thresh) / iou_thresh scores = torch.rand(N) return boxes, scores @@ -399,7 +403,12 @@ def test_nms_cuda(self): r_cpu = ops.nms(boxes, scores, iou) r_cuda = ops.nms(boxes.cuda(), scores.cuda(), iou) - self.assertTrue(torch.allclose(r_cpu, r_cuda.cpu()), err_msg.format(iou)) + is_eq = torch.allclose(r_cpu, r_cuda.cpu()) + if not is_eq: + # if the indices are not the same, ensure that it's because the scores + # are duplicate + is_eq = torch.allclose(scores[r_cpu], scores[r_cuda.cpu()]) + self.assertTrue(is_eq, err_msg.format(iou)) class NewEmptyTensorTester(unittest.TestCase):