From 2242bdf0a32e49d6692e8a5e70677f110d81e61d Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sat, 24 Feb 2024 22:30:14 +1100
Subject: [PATCH] Update gemma.py

---
 unsloth/models/gemma.py | 46 ++++++++++++++++++++---------------------
 1 file changed, 23 insertions(+), 23 deletions(-)

diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 243c9d9d7b..f32c10884c 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -467,30 +467,30 @@ def GemmaForCausalLM_fast_forward(
     pass
 
     loss = None
-    # if labels is not None:
-    #     shift_logits = logits
-    #     if not hasattr(self, "extra_ignored_labels"):
-    #         # Fixes https://github.com/unslothai/unsloth/issues/10
-    #         self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda")
-    #     pass
-        
-    #     shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
-    #     loss = fast_cross_entropy_loss(
-    #         logits = shift_logits,
-    #         labels = shift_labels,
-    #     )
-    # pass
     if labels is not None:
-        # Shift so that tokens < n predict n
-        shift_logits = logits[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-        # Flatten the tokens
-        loss_fct = torch.nn.CrossEntropyLoss()
-        shift_logits = shift_logits.view(-1, self.config.vocab_size)
-        shift_labels = shift_labels.view(-1)
-        # Enable model parallelism
-        shift_labels = shift_labels.to(shift_logits.device)
-        loss = loss_fct(shift_logits, shift_labels)
+        shift_logits = logits
+        if not hasattr(self, "extra_ignored_labels"):
+            # Fixes https://github.com/unslothai/unsloth/issues/10
+            self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda")
+        pass
+        
+        shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
+        loss = fast_cross_entropy_loss(
+            logits = shift_logits,
+            labels = shift_labels,
+        )
+    pass
+    # if labels is not None:
+    #     # Shift so that tokens < n predict n
+    #     shift_logits = logits[..., :-1, :].contiguous()
+    #     shift_labels = labels[..., 1:].contiguous()
+    #     # Flatten the tokens
+    #     loss_fct = torch.nn.CrossEntropyLoss()
+    #     shift_logits = shift_logits.view(-1, self.config.vocab_size)
+    #     shift_labels = shift_labels.view(-1)
+    #     # Enable model parallelism
+    #     shift_labels = shift_labels.to(shift_logits.device)
+    #     loss = loss_fct(shift_logits, shift_labels)
 
     if not return_dict:
         output = (logits,) + outputs[1:]