From 63f41575f4f70989891c90fc2e1ff9bd0451d4d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 21 Nov 2024 04:10:24 -0800 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5cefab90e1..e741c1c628 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -53,6 +53,10 @@ def _wrap_fast_inference(generate, device_type, dtype, model): kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) + if "pixel_values" in kwargs: + kwargs["pixel_values"] = kwargs["pixel_values"].to(model.dtype) + pass + # Autocasted with torch.autocast(device_type = device_type, dtype = dtype): output = generate(*args, **kwargs)