Skip to content

Commit

Permalink
Merge pull request #6 from pcuenca/add_palma_shift_mask
Browse files Browse the repository at this point in the history
Shift mask from `1:`
  • Loading branch information
molbap committed May 13, 2024
2 parents 404abd8 + 60ad9c5 commit daccbe7
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 daccbe7

Please sign in to comment.