From 63642327576fdae2fbd8159d68a18e643449cfc4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 2 Feb 2025 03:57:24 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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")