| """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" | |