diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 401b8986a8..d0ffa53d5b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -937,7 +937,7 @@ def LlamaModel_fast_forward_inference(
     next_decoder_cache = []
     residual = torch.empty_like(X)
     print(bsz, q_len, hd)
-    _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32)
+    _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
     XX, XX2 = _XX[0], _XX[1]
     variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
     temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")