Add `int` type to `device` parameter of torch.set_default_device() on the doc

[hyperkai]

📚 The doc issue

The doc of torch.set_default_device() explains device parameter but it doesn't have int type as shown below:

device (device or string) – the device to set as default

Actually, device parameter with int type works as shown below:

import torch

torch.set_default_device(device=0) # Here

torch.tensor([0, 1, 2]).device
# device(type='cuda', index=0)

Suggest a potential alternative/fix

The doc should also say int type as shown below:

device (device, string or int) – the device to set as default

cc @svekars @brycebortree @albanD

[mshr-h]

I agree. But according to the official PyTorch docs, it's a legacy behavior.
Tensor Attributes — PyTorch 2.3 documentation

For legacy reasons, a device can be constructed via a single device ordinal, which is treated as a cuda device. This matches Tensor.get_device(), which returns an ordinal for cuda tensors and is not supported for cpu tensors.

torch.device(1)
device(type='cuda', index=1)

I think we also need to add an example of something like this.

>>> torch.set_default_device(0) # Legacy
>>> torch.get_default_device()
device(type='cuda', index=0)

[malfet]

Not sure why this has a triage review to be frank. Feels like a doc issue and there is already a PR attempting to fix it

[albanD]

The 3 options we have here are:

• Keep the current behavior of meaning the index of the cuda device
• Remove this option everywhere
• Change it to mean the index of the accelerator device

From the discussion during triage, even though passing an int here is confusing, I think we should go with the third one as it is backward compatible (this BC-breaking change has little benefit really) and makes sense in our device-generic world.

After this change, we should indeed fix typing and docs to reflect the expected behavior.