Skip to content

Commit 6ebde0f

Browse files
masahiylc
authored andcommitted
Fix flaky NMS test by making sure scores are unique (apache#9140)
1 parent 4e4874f commit 6ebde0f

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

tests/python/frontend/pytorch/test_forward.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1963,8 +1963,9 @@ def _gen_rand_inputs(num_boxes):
19631963
boxes = torch.rand(num_boxes, box_len, dtype=torch.float) * 0.5
19641964
boxes[:, 2] += boxes[:, 0]
19651965
boxes[:, 3] += boxes[:, 1]
1966-
scores = torch.from_numpy(np.random.uniform(-1, 1, size=(num_boxes,)).astype(np.float32))
1967-
return boxes, scores
1966+
scores = np.linspace(0, 1, num=num_boxes).astype("float32")
1967+
np.random.shuffle(scores)
1968+
return boxes, torch.from_numpy(scores)
19681969

19691970
targets = ["llvm", "cuda"]
19701971

0 commit comments

Comments
 (0)