File size: 558 Bytes
580cc4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | """Shared device selection: cuda > mps > cpu.
Keeping the priority order in one place means the Apple Silicon path never
leaks platform-specific branches into the training and eval scripts - they all
just ask for the best available device. AMP/GradScaler stay CUDA-only at the
call sites that use them; MPS runs the same code in fp32.
"""
import torch
def pick_device(prefer=None):
if prefer:
return prefer
if torch.cuda.is_available():
return "cuda"
if torch.backends.mps.is_available():
return "mps"
return "cpu"
|