diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 4989f30db5..3e15b82876 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -284,13 +284,17 @@ pass
 def fast_mlp_inference(self, X):
     # gate = self.gate_proj(X)
     # up   = self.up_proj(X)
-    gate = fast_linear_forward(self.gate_proj, X)
-    up   = fast_linear_forward(self.  up_proj, X)
+    bsz, _, hd = X.shape
+    mlp_size = model.config.intermediate_size
+    # temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda")
+
+    gate = fast_linear_forward(self.gate_proj, X)#, out = temp[0])
+    up   = fast_linear_forward(self.  up_proj, X)#, out = temp[1])
     gate = torch.nn.functional.silu(gate, inplace = True)
     gate *= up
 
     # X = self.down_proj(gate)
-    down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:X.shape[2]])
+    down = fast_linear_forward(self.down_proj, gate)#, out = up[:,:,:hd])
     return down
 pass