From d5becd080e6101ec8f61fd0ebcf98a997f110e04 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 22 Oct 2024 00:29:30 -0700
Subject: [PATCH] Update mistral.py

---
 unsloth/models/mistral.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 15e9efc426..00dcc5cd1d 100644
--- 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