From 88060b94f1a7f22b8dfafcfc3f23d76f910797ae Mon Sep 17 00:00:00 2001 From: sandel Date: Fri, 18 Sep 2026 02:13:13 +1000 Subject: [PATCH] Ignore padding positions in the replaced-token detection loss Signed-off-by: sandel <326855337+gss10282025@users.noreply.github.com> Assisted-by: OpenAI Codex --- diffcse/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffcse/models.py b/diffcse/models.py index 069d5b1..6d3e091 100644 --- a/diffcse/models.py +++ b/diffcse/models.py @@ -269,7 +269,7 @@ def cl_forward(cls, cls.electra_fix_acc = float(1.0 - (prediction*fix).sum()/fix.sum()) cls.electra_acc = float(((prediction == e_labels) * attention_mask).sum()/attention_mask.sum()) # masked_lm_loss = loss_fct(prediction_scores.view(-1, cls.config.vocab_size), mlm_labels.view(-1)) - masked_lm_loss = loss_fct(prediction_scores.view(-1, 2), e_labels.view(-1)) + masked_lm_loss = loss_fct(prediction_scores.view(-1, 2), e_labels.masked_fill(attention_mask == 0, -100).view(-1)) loss = loss + cls.model_args.lambda_weight * masked_lm_loss if not return_dict: