Update vision.py

2025-03-14 05:41:17 -07:00 · 2025-03-14 05:41:17 -07:00 · f988ed4851
commit f988ed4851
parent 26045d8efe
1 changed files with 2 additions and 0 deletions
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@ -102,6 +102,8 @@ def unsloth_base_fast_generate(
        global NUM_LOGITS_TO_KEEP
        if arch not in NUM_LOGITS_TO_KEEP:
            m = self
+            # Find which is needed ie
+            # num_logits_to_keep or logits_to_keep
            while hasattr(m, "model"):
                if hasattr(m, "forward"):
                    keys = inspect.signature(m.forward).parameters.keys()