use torch.testing.assert_close

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Lightning-AI · Aug 22, 2024 · 941b132 · 941b132
1 parent c92e8a8
commit 941b132
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/thunder/tests/test_inplace_functionalization.py b/thunder/tests/test_inplace_functionalization.py
@@ -580,9 +580,9 @@ def f(x, y, idx, src):
     jitted = executor.make_callable(f)
     o = jitted(a, b, idx, src)
 
-    assert a.allclose(a_)
-    assert b.allclose(b_)
-    assert o.allclose(o_)
+    torch.testing.assert_close(a, a_)
+    torch.testing.assert_close(b, b_)
+    torch.testing.assert_close(o, o_)
 
 
 @instantiate(