From bef3496dbdf6216eaa3fc8a1bcd3b9a3da30fd6b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 21 Nov 2024 04:05:05 -0800 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 87a3846fdb..5cefab90e1 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -54,7 +54,7 @@ def _wrap_fast_inference(generate, device_type, dtype, model): kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) # Autocasted - with torch.autocast(device_type = model.device.type, dtype = dtype): + with torch.autocast(device_type = device_type, dtype = dtype): output = generate(*args, **kwargs) pass return output @@ -341,6 +341,7 @@ class FastBaseVisionModel: if dtype == "float16": dtype = torch.float16 elif dtype == "bfloat16": dtype = torch.bfloat16 pass + device_type = model.device.type # Wrap model.generate if model.generate.__name__ != "_fast_generate":