diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 9e53176d52..acf999faff 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -126,8 +126,8 @@ def unsloth_base_fast_generate(
         pass
     pass
     key = NUM_LOGITS_TO_KEEP[arch]
-    # if key is not None and key not in kwargs:
-    #     kwargs[key] = 1
+    if key is not None and key not in kwargs:
+        kwargs[key] = 1
     global PROMPT_LOOPKUP
     if arch not in PROMPT_LOOPKUP:
         PROMPT_LOOPKUP[arch] = True
@@ -146,8 +146,6 @@ def unsloth_base_fast_generate(
     try: kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype)
     except: pass
 
-    print(args, kwargs, self._old_generate)
-
     # Mixed precision autocast
     if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
         autocaster = torch.autocast(device_type = "cuda", dtype = dtype)