[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/spec_decode/utils.py
+++ b/tests/spec_decode/utils.py
@@ -161,7 +161,7 @@ def assert_logprobs_dict_allclose(
                single_step_actual_logprobs[token_id].logprob)
            expected = torch.tensor(
                single_step_expected_logprobs[token_id].logprob)
-            assert torch.allclose(actual, expected)
+            torch.testing.assert_close(actual, expected)


 def create_sampler_output_list(