From 2f5cc5393eb28413fca14b2e202ff762a5199c5e Mon Sep 17 00:00:00 2001 From: Vasilis Vryniotis Date: Thu, 11 Nov 2021 13:35:34 +0000 Subject: [PATCH] Set seed on test_nms_ref to reduce flakiness --- test/test_ops.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index cd7f992a7c8..c8e4e396c7e 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -479,7 +479,9 @@ def _create_tensors_with_iou(self, N, iou_thresh): return boxes, scores @pytest.mark.parametrize("iou", (0.2, 0.5, 0.8)) - def test_nms_ref(self, iou): + @pytest.mark.parametrize("seed", range(10)) + def test_nms_ref(self, iou, seed): + torch.random.manual_seed(seed) err_msg = "NMS incompatible between CPU and reference implementation for IoU={}" boxes, scores = self._create_tensors_with_iou(1000, iou) keep_ref = self._reference_nms(boxes, scores, iou)