torch.cuda.set_device cannot use to set cpu device, but give an ambiguity hint

[rubbberrabbit]

🐛 Describe the bug

when give the string parameter as input to set the device torch use，the hint shows that cpu is allow

import torch
torch.cuda.set_device("gpu:0") # it should be torch.cuda.set_device("cuda:0") but this is to get the hint
# RuntimeError: Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: gpu

but if I follow the hint, and try to use torch.cuda.set_device to set cpu as target device

import torch
torch.cuda.set_device("cpu")
#  ValueError: Expected a cuda device, but got: cpu
device = torch.device("cpu")
torch.cuda.set_device(device)
# ValueError: Expected a cuda device, but got: cpu

in neither way, the api won't allow me to do so, so the hint is pretty ambiguity.
I went to look at the torch.cuda.set_device source code and found
the torch.cuda.set_device use _get_device_index and the default parameter is

def _get_device_index(device: Any, optional: bool = False,
                      allow_cpu: bool = False) -> int:
    ...

and actually won't allow the cpu set.
if device = _get_device_index(device,allow_cpu = True) is more reasonable in the torch.cuda.set_device call, considering the hint and user can use the function to set the global device as cpu.

Versions

pytorch 2.1.0 and the latest version

[drisspg]

torch.cuda.device() is used for specifying which "cuda" device you are trying to access, specifically when you have multiple visible cuda devices in your environment. What would you expect to happen by setting this to CPU?

[rubbberrabbit]

I see, I want to use CPU device as a small batch experiment, I want to make sure the whole process will use CPU, CUDA_VISIBLE_DEVICES = -1 is good to me. I think the hint is a liite ambiguity as it said
"Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia" but wont work with cpu