Update vision.py

This commit is contained in:
Daniel Han 2024-11-21 04:01:17 -08:00
commit 2b8b2f159f

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 = device_type, dtype = dtype):
with torch.autocast(device_type = model.device.type, dtype = dtype):
output = generate(*args, **kwargs)
pass
return output