Skip to content

Commit 9a8b215

Browse files
julber95meilame-tayebjee
authored andcommitted
fix: invert attention mask in LabelAttention to correctly ignore padding tokens
1 parent f269062 commit 9a8b215

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

‎torchTextClassifiers/model/components/text_embedder.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,7 @@ def forward(
279279
attn_mask = None
280280
if attention_mask is not None:
281281
# Convert: 0 (padding) -> True (mask out), 1 (real) -> False (attend to)
282-
attn_mask = attention_mask == 0 # (B, T)
282+
attn_mask = attention_mask == 1 # (B, T)
283283
# Expand to (B, 1, 1, T) for broadcasting across heads and queries
284284
attn_mask = attn_mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, T)
285285

0 commit comments

Comments
 (0)