mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
12 lines
266 B
Python
12 lines
266 B
Python
|
import torch
|
||
|
|
|
||
|
|
def choose_torch_device() -> str:
|
||
|
|
'''Convenience routine for guessing which GPU device to run model on'''
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
return 'cuda'
|
||
|
|
if torch.backends.mps.is_available():
|
||
|
|
return 'mps'
|
||
|
|
return 'cpu'
|
||
|
|
|
||
|
|
|