From 9df19aeb9fec1e985d4e9f8d2b273f25baa1a0d8 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 2 Feb 2024 20:32:21 +1100 Subject: [PATCH] past_key_values --- unsloth/models/llama.py | 2 +- unsloth/models/mistral.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index fa3c8d1c5a..052ff54282 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -750,7 +750,7 @@ def LlamaForCausalLM_fast_forward( # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) self.model._has_no_labels = labels is None - if past_key_value is not None and \ + if past_key_values is not None and \ hasattr(self.model.layers[0].self_attn, "paged_attention"): outputs = LlamaModel_fast_forward_inference( self.model, diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 99dc3c1fd9..cd92d0e7dd 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -201,7 +201,7 @@ def MistralForCausalLM_fast_forward( # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) self.model._has_no_labels = labels is None - if past_key_value is not None and \ + if past_key_values is not None and \ hasattr(self.model.layers[0].self_attn, "paged_attention"): outputs = LlamaModel_fast_forward_inference( self.model,