added simple test

be87cc98 · Duc · 2e57b128 · be87cc98
Commit be87cc98 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
@@ -21,6 +21,9 @@ def test_fps(dtype, device):
    ], dtype, device)
    batch = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)

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