From bc1746839c30e8c724dc5bc3b48693e53b20fefe Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 18 Mar 2024 03:39:44 +1100
Subject: [PATCH] dtype

---
 unsloth/models/llama.py   | 3 ++-
 unsloth/models/mistral.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index b5fd37056c..262649f7fc 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -509,6 +509,8 @@ def LlamaModel_fast_forward(
     if inputs_embeds is None:
         inputs_embeds = self.embed_tokens(input_ids)
 
+    inputs_embeds = inputs_embeds.to(self.config.torch_dtype)
+
     # Normalized from Gemma
     IS_GEMMA = self.config.model_type == "gemma"
     train_embed_tokens = self.embed_tokens.weight.requires_grad
@@ -756,7 +758,6 @@ def CausalLM_fast_forward(fast_forward_inference):
         else:
             logits = self.lm_head(hidden_states)
         pass
-        print(logits.dtype)
 
         loss = None
         if labels is not None:
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index c1e39e4a2e..e959b731a8 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -230,6 +230,7 @@ def MistralForCausalLM_fast_forward(
     else:
         logits = self.lm_head(hidden_states)
     pass
+    print(logits.dtype)
 
     loss = None
     if labels is not None: