Update mistral.py

2024-10-22 00:29:30 -07:00 · 2024-10-22 00:29:30 -07:00 · d5becd080e
commit d5becd080e
parent 1ca9a03ab7
1 changed files with 3 additions and 2 deletions
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@ -254,8 +254,9 @@ def MistralForCausalLM_fast_forward(
        
        shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
        loss = fast_cross_entropy_loss(
-            logits = shift_logits,
-            labels = shift_labels,
+            logits  = shift_logits,
+            labels  = shift_labels,
+            n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None),
        )
    pass