Skip to content

Commit

Permalink
Shift attention mask from 1:
Browse files Browse the repository at this point in the history
After discussion with @molbap
  • Loading branch information
pcuenca committed May 13, 2024
1 parent 404abd8 commit 60ad9c5
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/transformers/models/paligemma/modeling_paligemma.py
Original file line number Diff line number Diff line change
Expand Up @@ -463,7 +463,7 @@ def forward(
if attention_mask.dim() == 4:
# take top or bottom row of the 4d mask.
# this should only be used in the initial pass with full attention on prefix.
shift_attention_mask = attention_mask[:, 0, 0, :-1].squeeze(1) if not left_padding else attention_mask[:, 0, -1, :-1].squeeze(1)
shift_attention_mask = attention_mask[:, 0, 0, 1:].squeeze(1) if not left_padding else attention_mask[:, 0, -1, 1:].squeeze(1)
elif attention_mask.dim() == 2:
# take normal slice of the attn mask
shift_attention_mask = attention_mask[..., 1:]
Expand Down

0 comments on commit 60ad9c5

Please sign in to comment.