fix mask shift

Author: jmercat

open_lm/utils/transformers/hf_model.py

@@ -140,7 +140,7 @@ 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)
                 shift_mask = torch.logical_and(shift_mask.view(-1), shift_labels != -100)