Update vision.py
This commit is contained in:
parent
3d22829736
commit
2b8b2f159f
1 changed files with 1 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue