From 5ba087859adc98c250cebf29bab509ebf899f8da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 21 Mar 2025 17:26:20 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 1 + 1 file changed, 1 insertion(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 0a43f46cbf..cf3d089d3c 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -102,6 +102,7 @@ def _fast_prepare_inputs_for_generation(self, input_ids, **kwargs,): if "past_key_values" in kwargs: print("FIX", input_ids.shape) input_ids = input_ids[:,[-1]] + print("FIX AFTER", input_ids.shape) kwargs["attention_mask"] = kwargs["attention_mask"][:,[-1]] if "cache_position" in kwargs: kwargs["position_ids"] = kwargs["cache_position"]