ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2025-03-21 17:26:20 -07:00
commit 5ba087859a
1 changed files with 1 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/llama.py
View file

@ -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"]