nukki / app /device.py
minsuas's picture
Upload folder using huggingface_hub
bcfaab3 verified
Raw
History Blame Contribute Delete
620 Bytes
"""디바이스 선택 (Apple Silicon MPS 우선)."""
import torch
_DEVICE = None
def get_device() -> torch.device:
global _DEVICE
if _DEVICE is not None:
return _DEVICE
if torch.backends.mps.is_available():
_DEVICE = torch.device("mps")
elif torch.cuda.is_available():
_DEVICE = torch.device("cuda")
else:
_DEVICE = torch.device("cpu")
return _DEVICE
def device_name() -> str:
d = get_device()
if d.type == "mps":
return "Apple Silicon (MPS)"
if d.type == "cuda":
return f"CUDA ({torch.cuda.get_device_name(0)})"
return "CPU"