Skip to content

Commit 2945515

Browse files
Merge pull request #5 from liulog/fix-self_attention-test
fix test/ops/self_attention.py
2 parents d44b46b + 896616b commit 2945515

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

test/ops/self_attention.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def torch_self_attention(attn_val, query, key, value, scale):
1515
L, S = query.size(-2), key.size(-2)
1616
attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
1717

18-
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
18+
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=S-L)
1919
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
2020
attn_bias.to(query.dtype)
2121

0 commit comments

Comments
 (0)