torch.set_default_device β PyTorch 2.7 documentation (original) (raw)
torch.set_default_device(device)[source][source]ΒΆ
Sets the default torch.Tensor
to be allocated on device
. This does not affect factory function calls which are called with an explicitdevice
argument. Factory calls will be performed as if they were passed device
as an argument.
To only temporarily change the default device instead of setting it globally, use with torch.device(device):
instead.
The default device is initially cpu
. If you set the default tensor device to another device (e.g., cuda
) without a device index, tensors will be allocated on whatever the current device for the device type, even after torch.cuda.set_device() is called.
Parameters
device (device or string) β the device to set as default
Example:
torch.get_default_device() device(type='cpu') torch.set_default_device('cuda') # current device is 0 torch.get_default_device() device(type='cuda', index=0) torch.set_default_device('cuda') torch.cuda.set_device('cuda:1') # current device is 1 torch.get_default_device() device(type='cuda', index=1) torch.set_default_device('cuda:1') torch.get_default_device() device(type='cuda', index=1)