Relax some numeric checks to accommodate Windows machine CI

nod-ai · Sep 25, 2024 · 8a24c5a · 8a24c5a
1 parent cfa705a
commit 8a24c5a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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
         )