| 88 | |
| 89 | |
| 90 | def normalize_device(backend: str, device_index: int) -> str: |
| 91 | if backend == "cuda": |
| 92 | if not torch.cuda.is_available(): |
| 93 | raise RuntimeError("CUDA backend requested but torch.cuda.is_available() is false") |
| 94 | torch.cuda.set_device(device_index) |
| 95 | return "cuda" |
| 96 | if backend == "mps": |
| 97 | if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): |
| 98 | raise RuntimeError("MPS backend requested but torch.backends.mps.is_available() is false") |
| 99 | return "mps" |
| 100 | if backend == "xpu": |
| 101 | if not (hasattr(torch, "xpu") and torch.xpu.is_available()): |
| 102 | raise RuntimeError("XPU backend requested but torch.xpu.is_available() is false") |
| 103 | return "xpu" |
| 104 | return "cpu" |
| 105 | |
| 106 | |
| 107 | def sync_device(device: str) -> None: |