From 4bdb66880994892f1a483e318dbddf0b7e5f7b7f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Mar 2025 23:34:52 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e415e50cdc..a96b435b58 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1562,7 +1562,9 @@ def unsloth_fast_generate( # For newer HF kwargs["cache_implementation"] = "dynamic" # For num_logits_to_keep - if "num_logits_to_keep" not in kwargs or "logits_to_keep" not in kwargs: + num_logits_to_keep = kwargs.get("num_logits_to_keep", None) + logits_to_keep = kwargs.get("logits_to_keep", None) + if num_logits_to_keep is None and logits_to_keep is None: kwargs["num_logits_to_keep"] = 1 # Remove token_type_ids