r""" Resolve the device based on the acceleration options and the available devices in the system Rules: 1. AUTO: Check for the best available device on the system. 2. User-defined: Check if the device actually exists, otherwise fall-back to CPU
(accelerator_device: AcceleratorDevice)
| 8 | |
| 9 | |
| 10 | def decide_device(accelerator_device: AcceleratorDevice) -> str: |
| 11 | r""" |
| 12 | Resolve the device based on the acceleration options and the available devices in the system |
| 13 | Rules: |
| 14 | 1. AUTO: Check for the best available device on the system. |
| 15 | 2. User-defined: Check if the device actually exists, otherwise fall-back to CPU |
| 16 | """ |
| 17 | cuda_index = 0 |
| 18 | device = "cpu" |
| 19 | |
| 20 | has_cuda = torch.backends.cuda.is_built() and torch.cuda.is_available() |
| 21 | has_mps = torch.backends.mps.is_built() and torch.backends.mps.is_available() |
| 22 | |
| 23 | if accelerator_device == AcceleratorDevice.AUTO: |
| 24 | if has_cuda: |
| 25 | device = f"cuda:{cuda_index}" |
| 26 | elif has_mps: |
| 27 | device = "mps" |
| 28 | |
| 29 | else: |
| 30 | if accelerator_device == AcceleratorDevice.CUDA: |
| 31 | if has_cuda: |
| 32 | device = f"cuda:{cuda_index}" |
| 33 | else: |
| 34 | _log.warning("CUDA is not available in the system. Fall back to 'CPU'") |
| 35 | elif accelerator_device == AcceleratorDevice.MPS: |
| 36 | if has_mps: |
| 37 | device = "mps" |
| 38 | else: |
| 39 | _log.warning("MPS is not available in the system. Fall back to 'CPU'") |
| 40 | |
| 41 | _log.info("Accelerator device: '%s'", device) |
| 42 | return device |