set_gpu_device#
- ptyrad.utils.common.set_gpu_device(gpuid=0)[source]#
Sets the GPU device based on the input. If ‘acc’ is passed, it returns None to defer to accelerate.
- Parameters:
gpuid (str or int) – The GPU ID to use. Can be: - “acc”: Defer device assignment to accelerate. - “cpu”: Use CPU. - An integer (or string representation of an integer) for a specific GPU ID. This only has effect on NVIDIA GPUs.
- Returns:
The selected device, or None if deferred to accelerate.
- Return type:
torch.device or None