added a test for vector ratio

5f1939fd · Duc · 01a51d47 · 5f1939fd
Commit 5f1939fd authored Dec 05, 2020 by Duc
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

test/test_fps.py test/test_fps.py +3 -0

No files found.
--- a/test/test_fps.py
+++ b/test/test_fps.py
@@ -27,6 +27,9 @@ def test_fps(dtype, device):
    out = fps(x, batch, ratio=torch.tensor(0.5), random_start=False)
    assert out.tolist() == [0, 2, 4, 6]

+    out = fps(x, batch, ratio=torch.tensor([0.5, 0.5]), random_start=False)
+    assert out.tolist() == [0, 2, 4, 6]
+
    out = fps(x, ratio=torch.tensor(0.5), random_start=False)
    assert out.sort()[0].tolist() == [0, 5, 6, 7]