()
| 87 | # ────────────────────────────────────────────────────────────────────────────── |
| 88 | |
| 89 | def get_device() -> torch.device: |
| 90 | if torch.cuda.is_available(): |
| 91 | return torch.device("cuda") |
| 92 | if torch.backends.mps.is_available(): |
| 93 | return torch.device("mps") |
| 94 | return torch.device("cpu") |
| 95 | |
| 96 | |
| 97 | def get_autocast_ctx(device: torch.device): |