tolerance adjustment

nod-ai · Jan 9, 2025 · 8d8d2ad · 8d8d2ad
1 parent ae6e0ca
commit 8d8d2ad
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/sharktank/tests/kernels/attention_template_test.py b/sharktank/tests/kernels/attention_template_test.py
@@ -18,14 +18,14 @@
 from sharktank.types import layout_utils
 
 
-class punet_attention(unittest.TestCase):
+class custom_attention(unittest.TestCase):
     def setUp(self):
         torch.manual_seed(420)
 
     @parameterized.expand(
         [
-            (torch.float32, 1e-3, 1e-4),
-            (torch.float16, 1e-3, 1e-4),
+            (torch.float32, 5e-3, 1e-3),
+            (torch.float16, 5e-3, 1e-3),
         ]
     )
     def test_compare_torch_spda(self, dtype, atol, rtol):