From 7b28da7487b534f1045852aee37f9f9817eed5d8 Mon Sep 17 00:00:00 2001 From: julber95 Date: Thu, 7 May 2026 14:56:40 +0000 Subject: [PATCH] fix: invert attention mask in LabelAttention to correctly ignore padding tokens --- torchTextClassifiers/model/components/text_embedder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchTextClassifiers/model/components/text_embedder.py b/torchTextClassifiers/model/components/text_embedder.py index be5e7fe..f36b32a 100644 --- a/torchTextClassifiers/model/components/text_embedder.py +++ b/torchTextClassifiers/model/components/text_embedder.py @@ -279,7 +279,7 @@ def forward( attn_mask = None if attention_mask is not None: # Convert: 0 (padding) -> True (mask out), 1 (real) -> False (attend to) - attn_mask = attention_mask == 0 # (B, T) + attn_mask = attention_mask == 1 # (B, T) # Expand to (B, 1, 1, T) for broadcasting across heads and queries attn_mask = attn_mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, T)