Skip to content

Fix/attention mask mismatch - #4218

Closed
KatharinaJacoby wants to merge 4 commits into
huggingface:mainfrom
KatharinaJacoby:fix/attention-mask-mismatch
Closed

KatharinaJacoby wants to merge 4 commits into
huggingface:mainfrom
KatharinaJacoby:fix/attention-mask-mismatch

Conversation

@KatharinaJacoby

Copy link
Copy Markdown

The issue seems to come from the concatenation of the virtual attention mask, which changes the dimensions of the attention_mask and causes a mismatch with per_token_entropy.
Checking that the dimensions of per_token_entropy and attention_mask match before multiplication. I changed the code block so the attention_mask is truncated to match the sequence length of per_token_entropy.
And added a unit test:
-set up common test data in the setUp method.
-test the scenario where the attention mask and per-token entropy have matching dimensions.
-test the scenario where the attention mask and per-token entropy have mismatched dimensions.
-test the scenario where position_ids is used instead of attention_mask.
-test the scenario where neither attention_mask nor position_ids is present, expecting a ValueError.

attention_mask is truncated to match the sequence length of per_token_entropy. This should fix the tensor mismatch error.
@pramodith

Copy link
Copy Markdown
Collaborator

Hey @KatharinaJacoby I think another PR that I opened last week resolves the issue that your PR covers too. #4196

@qgallouedec

Copy link
Copy Markdown
Member

closing in favor of #4196, thanks a lot @KatharinaJacoby and @pramodith

@qgallouedec qgallouedec closed this Oct 9, 2025
@KatharinaJacoby
KatharinaJacoby deleted the fix/attention-mask-mismatch branch October 10, 2025 06:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants