(device_spec: Union[str, int, List[int]])
| 185 | |
| 186 | |
| 187 | def get_device(device_spec: Union[str, int, List[int]]) -> torch.device: |
| 188 | ddp_flag = False |
| 189 | if isinstance(device_spec, (list, ListConfig)): |
| 190 | ddp_flag = True |
| 191 | device_spec = initialize_distributed() |
| 192 | if torch.cuda.is_available() and "cuda" in str(device_spec): |
| 193 | return torch.device(device_spec), ddp_flag |
| 194 | if not torch.cuda.is_available(): |
| 195 | if device_spec != "cpu": |
| 196 | logger.warning(f"❎ Device spec: {device_spec} not support, Choosing CPU instead") |
| 197 | return torch.device("cpu"), False |
| 198 | |
| 199 | device = torch.device(device_spec) |
| 200 | return device, ddp_flag |
| 201 | |
| 202 | |
| 203 | class PostProcess: |
no test coverage detected