diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c91f04073c..d1d5f5e16e 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -920,6 +920,11 @@ def LlamaModel_fast_forward_inference(
     bsz, q_len = input_ids.shape
     hd = self.config.hidden_size
     mlp_size = self.config.intermediate_size
+
+    X = self.model.embed_tokens(input_ids)
+    X = X.to(self.config.torch_dtype)
+    bsz, q_len, hd = X.shape
+    assert(q_len == 1)
     
     # Get saved buffers to reduce memory movement
     residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
@@ -929,11 +934,6 @@ def LlamaModel_fast_forward_inference(
     temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
     temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
 
-    X = self.model.embed_tokens(input_ids)
-    X = X.to(self.config.torch_dtype)
-    bsz, q_len, hd = X.shape
-    assert(q_len == 1)
-
     seq_len = past_key_values[0][0].shape[-2]
     if bsz != 1:
         attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(