lint

mosaicml · Jan 13, 2024 · 4845113 · 4845113
1 parent 054cf15
commit 4845113
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/models/layers/test_flash_triton_torch.py b/tests/models/layers/test_flash_triton_torch.py
@@ -261,7 +261,9 @@ def gen_bias(attn_impl: str):
             'rope_impl'] == 'hf'
 
         # special case that (likely) fails due to numerics
-        if clip_qkv and (qk_ln or qk_gn) and using_hf_rope and attn_type == 'grouped_query_attention':
+        if clip_qkv and (
+                qk_ln or qk_gn
+        ) and using_hf_rope and attn_type == 'grouped_query_attention':
             assert allclose_helper(p.grad, tp.grad, atol=2.e-2, rtol=2.e-2)
         else:
             assert allclose_helper(p.grad, tp.grad)