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)