Update vision.py

2024-11-21 04:10:43 -08:00 · 2024-11-21 04:10:43 -08:00 · a98fc9c57c
commit a98fc9c57c
parent 5886ecb246
1 changed files with 3 additions and 2 deletions
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@ -53,9 +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:
+        try:
            kwargs["pixel_values"] = kwargs["pixel_values"].to(model.dtype)
-        pass
+        except:
+            pass

        # Autocasted
        with torch.autocast(device_type = device_type, dtype = dtype):