ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update vision.py

Browse source
This commit is contained in:
Daniel Han 2024-11-21 04:05:05 -08:00
commit bd0c5bb1a1
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/vision.py
View file

@ -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":