Return the best available device, prioritising MPS for Apple silicon.
(preferred: str | None = None)
| 26 | |
| 27 | |
| 28 | def resolve_device(preferred: str | None = None) -> torch.device: |
| 29 | """Return the best available device, prioritising MPS for Apple silicon.""" |
| 30 | if preferred: |
| 31 | if preferred == "mps" and torch.backends.mps.is_available(): |
| 32 | return torch.device("mps") |
| 33 | if preferred == "cuda" and torch.cuda.is_available(): |
| 34 | return torch.device("cuda") |
| 35 | if preferred == "cpu": |
| 36 | return torch.device("cpu") |
| 37 | if torch.backends.mps.is_available(): |
| 38 | return torch.device("mps") |
| 39 | if torch.cuda.is_available(): |
| 40 | return torch.device("cuda") |
| 41 | return torch.device("cpu") |
| 42 | |
| 43 | |
| 44 | class SmallMLP(nn.Module): |