use the wrapped function, and check for success

mberr · Feb 7, 2024 · 45e5c15 · 45e5c15
1 parent 9949725
commit 45e5c15
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_decorator.py b/tests/test_decorator.py
@@ -171,5 +171,5 @@ def test_large():
 
     x = torch.rand(100000, 100, device="mps")
     y = torch.rand(200000, 100, device="mps")
-    z = knn(x, y, batch_size=x.shape[0])
-    assert z
+    _result, (batch_size,) = wrapped_knn(x, y, batch_size=x.shape[0])
+    assert batch_size > 0