From 6cc98350213795dbc30a735670153cb486f3fc61 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 2 Feb 2024 13:52:12 +1100 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 25ad680008..60fa43bc8e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -266,7 +266,7 @@ def LlamaAttention_fast_forward_inference( self.paged_attention_K = self.paged_attention[:,0] self.paged_attention_V = self.paged_attention[:,1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) - self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) + self.paged_attention_K[:seq_len] = V1.permute(2, 0, 1, 3) pass self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3) self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)