Small test fix to have right padding (#757)

mosaicml · Nov 22, 2023 · 1eecd4c · 1eecd4c
1 parent 1793c36
commit 1eecd4c
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 // 3] = 0
+        attention_mask[:, -s // 3:] = 0
 
     def gen_bias(attn_impl: str):
         causal = True