r""" Sets the device for PyTorch.
(device: Union[torch.device, str, int, None])
| 117 | |
| 118 | |
| 119 | def torch_set_device(device: Union[torch.device, str, int, None]) -> None: |
| 120 | r""" |
| 121 | Sets the device for PyTorch. |
| 122 | """ |
| 123 | if is_torch_npu_available(): |
| 124 | torch.npu.set_device(device) |
| 125 | elif is_torch_cuda_available(): |
| 126 | torch.cuda.set_device(device) |
| 127 | |
| 128 | |
| 129 | def manual_seed_all(seed) -> None: |