From 8a24c5ab7519d01fc7e21c4d06cce0087945a244 Mon Sep 17 00:00:00 2001 From: Boian Petkantchin Date: Wed, 25 Sep 2024 11:06:34 -0500 Subject: [PATCH] Relax some numeric checks to accommodate Windows machine CI --- sharktank/tests/models/llama/sharded_llama_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sharktank/tests/models/llama/sharded_llama_test.py b/sharktank/tests/models/llama/sharded_llama_test.py index 41671208e..9dd27c4f8 100644 --- a/sharktank/tests/models/llama/sharded_llama_test.py +++ b/sharktank/tests/models/llama/sharded_llama_test.py @@ -152,5 +152,5 @@ def testToyModelCompareToUnsharded(self): sharded_model.cache.paged.unflatten_page_table(decode_sharded_cache_state) ).flatten(start_dim=1) torch.testing.assert_close( - actual_decode_cache_state, expected_decode_cache_state + actual_decode_cache_state, expected_decode_cache_state, atol=1e-4, rtol=1e-4 )