fix minor bug

ROCm · Aug 8, 2024 · 182479c · 182479c
1 parent 9690410
commit 182479c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_flash_attn.py b/tests/test_flash_attn.py
@@ -2823,7 +2823,7 @@ def test_flash_attn_varlen_deterministic(seqlen_q, seqlen_k, swap_sq_sk, d, caus
     )
 
     g = torch.randn_like(out)
-    test_backward = test_backward and (d <= MAX_HEADDIM_SM8x or d > 224) or (is_sm80 or is_sm90)
+    test_backward = test_backward and ((d <= MAX_HEADDIM_SM8x or d > 224) or (is_sm80 or is_sm90))
     if test_backward:
         dq, dk, dv = torch.autograd.grad(out, (q_unpad, k_unpad, v_unpad), g, retain_graph=True)
         for _ in range(50):