From 5866b08f71daaf25cf0fdb09ba77abb8aba59624 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Fri, 23 Feb 2024 04:11:08 +1100
Subject: [PATCH] gemma

---
 unsloth/kernels/cross_entropy_loss.py | 2 +-
 unsloth/models/loader.py              | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 81a2312b75..0a73a393ec 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -168,7 +168,7 @@ def fast_cross_entropy_loss(logits, labels):
             "Unsloth's Github page if you want a faster and more memory efficient kernel!"
         )
         loss = slow_cross_entropy_loss(
-            logits.view(batch*seq_len, d), # Must cast to float32 for numerical stability
+            logits.float().view(batch*seq_len, d), # Must cast to float32 for numerical stability
             labels.view(-1),
         )
         return loss
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index add8e9dca6..67a59c850c 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -14,7 +14,6 @@
 
 from .llama import FastLlamaModel, logger
 from .mistral import FastMistralModel
-from .gemma import FastGemmaModel
 from transformers import AutoConfig
 from transformers import __version__ as transformers_version
 from peft import PeftConfig, PeftModel
@@ -26,6 +25,8 @@ major, minor = transformers_version.split(".")[:2]
 major, minor = int(major), int(minor)
 SUPPORTS_FOURBIT = (major > 4) or (major == 4 and minor >= 37)
 SUPPORTS_GEMMA   = (major > 4) or (major == 4 and minor >= 38)
+if SUPPORTS_GEMMA:
+    from .gemma import FastGemmaModel
 del major, minor