File size: 721 Bytes
473c3a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
from __future__ import annotations
from logging import getLogger
import torch
logger = getLogger(__name__)
def select_optimal_device(device: str | None) -> str:
"""
Guess what your optimal device should be based on backend availability.
If you pass a device, we just pass it through.
:param device: The device to use. If this is not None you get back what you passed.
:return: The selected device.
"""
if device is None:
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
logger.info(f"Automatically selected device: {device}")
return device
|