r""" Gets the current available device.
()
| 54 | |
| 55 | |
| 56 | def get_current_device() -> torch.device: |
| 57 | r""" |
| 58 | Gets the current available device. |
| 59 | """ |
| 60 | if is_torch_xpu_available(): |
| 61 | device = 'xpu:{}'.format(os.environ.get('LOCAL_RANK', '0')) |
| 62 | elif is_torch_npu_available(): |
| 63 | device = 'npu:{}'.format(os.environ.get('LOCAL_RANK', '0')) |
| 64 | elif is_torch_mps_available(): |
| 65 | device = 'mps:{}'.format(os.environ.get('LOCAL_RANK', '0')) |
| 66 | elif is_torch_cuda_available(): |
| 67 | device = 'cuda:{}'.format(os.environ.get('LOCAL_RANK', '0')) |
| 68 | else: |
| 69 | device = 'cpu' |
| 70 | |
| 71 | return torch.device(device) |
| 72 | |
| 73 | |
| 74 | def get_all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor: |