diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 8b86feb35e..c31cb3d84b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -753,7 +753,8 @@ def CausalLM_fast_forward(fast_forward_inference):
         hidden_states = outputs[0]
         bsz, q_len, hd = hidden_states.shape
         if bsz == 1 and q_len == 1:
-            logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
+            lm_head = self.lm_head.weight
+            logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
             logits = logits.unsqueeze(0).unsqueeze(0)
         else:
             logits = self.lm_head(hidden_states)
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 2db71f0288..f650d8f12e 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -225,7 +225,8 @@ def MistralForCausalLM_fast_forward(
     hidden_states = outputs[0]
     bsz, q_len, hd = hidden_states.shape
     if bsz == 1 and q_len == 1:
-        logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
+        lm_head = self.lm_head.weight
+        logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
         logits = logits.unsqueeze(0).unsqueeze(0)
     else:
         logits = self.lm_head(hidden_states)