mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Update to exclude GTX 16xx from autocast selection (#576)
* Update to exclude GTX 1660/1650 from autocast selection
This commit is contained in:
parent
e6011631a1
commit
f9b272a7b9
@ -13,9 +13,12 @@ def choose_torch_device() -> str:
|
||||
def choose_autocast_device(device):
|
||||
'''Returns an autocast compatible device from a torch device'''
|
||||
device_type = device.type # this returns 'mps' on M1
|
||||
# autocast only for cuda, but GTX 16xx have issues with it
|
||||
if device_type == 'cuda':
|
||||
return device_type,autocast
|
||||
elif device_type == 'cpu':
|
||||
return device_type,nullcontext
|
||||
device_name = torch.cuda.get_device_name()
|
||||
if 'GeForce GTX 1660' in device_name or 'GeForce GTX 1650' in device_name:
|
||||
return device_type,nullcontext
|
||||
else:
|
||||
return device_type,autocast
|
||||
else:
|
||||
return 'cpu',nullcontext
|
||||
|
Loading…
Reference in New Issue
Block a user