From caae7668e07449b650becf058a6d52597e8905e8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 16 May 2025 23:20:18 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 83bd22f625..f0085b0f69 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -89,15 +89,20 @@ def unsloth_base_fast_generate( *args, **kwargs, ): - print(args, kwargs) if len(args) != 0: input_ids = args[0] elif "input_ids" in kwargs: input_ids = kwargs["input_ids"] elif "input" in kwargs: input_ids = kwargs["input_ids"] + elif "inputs" in kwargs: + input_ids = kwargs["inputs"] else: - raise TypeError("Unsloth: You need to pass in input_ids to .generate!") + key = next(iter(kwargs.keys())) + if type(kwargs["key"]) is not torch.Tensor: + raise TypeError("Unsloth: You need to pass in input_ids to .generate!") + input_ids = kwargs[key] + pass assert(type(input_ids) is torch.Tensor) bsz = input_ids.shape[0]