Fix hf_model loss

mlfoundations · May 4, 2024 · b4003c7 · b4003c7
1 parent d85fbc8
commit b4003c7
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/open_lm/utils/transformers/hf_model.py b/open_lm/utils/transformers/hf_model.py
@@ -140,10 +140,11 @@ def forward(
             shift_logits = shift_logits.view(-1, shift_logits.size(-1))
             shift_labels = shift_labels.view(-1).to(shift_logits.device)
             if loss_mask is not None:
-                shift_mask = loss_mask[..., 1:].contiguous()
+                shift_mask = loss_mask[..., :-1].contiguous()
                 loss_fct = nn.CrossEntropyLoss(reduction="none")
                 loss = loss_fct(shift_logits, shift_labels)
-                loss = loss[shift_mask.view(-1)].sum() / shift_mask.sum()
+                shift_mask = torch.logical_and(shift_mask.view(-1), shift_labels != -100)
+                loss = loss[shift_mask.view(-1)].sum()/shift_mask.sum()
             else:
                 loss_fct = nn.CrossEntropyLoss()
                 loss = loss_fct(shift_logits, shift_labels)