mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-11-19 03:25:41 +00:00
16 lines
277 B
Python
16 lines
277 B
Python
from functools import lru_cache
|
|
|
|
import torch
|
|
|
|
|
|
@lru_cache()
|
|
def get_device() -> str:
|
|
"""Return the best torch backend available."""
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
|
|
if torch.backends.mps.is_available():
|
|
return "mps:0"
|
|
|
|
return "cpu"
|