[Misc/Testing] Use torch.testing.assert_close (#7324)

2024-08-15 21:24:04 -07:00
parent e165528778
commit 50b8d08dbd
25 changed files with 197 additions and 188 deletions
--- a/tests/kernels/utils.py
+++ b/tests/kernels/utils.py
@@ -924,5 +924,5 @@ def assert_actual_matches_ideal(test_params: PhaseTestParameters,
    * output_under_test: actually observed output value
    '''
    ideal_output = test_params.packed_qkvo.ideal_output
-    assert torch.allclose(ideal_output,
-                          output_under_test.view_as(ideal_output))
+    torch.testing.assert_close(ideal_output,
+                               output_under_test.view_as(ideal_output))