cleaner

mosaicml · Nov 22, 2023 · 0029362 · 0029362
1 parent 95d458b
commit 0029362
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_flash_triton_torch.py b/tests/test_flash_triton_torch.py
@@ -103,7 +103,7 @@ def test_attn_impl(attn_impl_0: str,
     if pad_attention_mask:
         # zero out the last third of the attention mask
         # to simulate padding
-        attention_mask[:, (s * 2) // 3:] = 0
+        attention_mask[:, -s // 3:] = 0
 
     def gen_bias(attn_impl: str):
         causal = True