torch.set_default_device

torch.set_default_device(device)

Sets the default torch.Tensor to be allocated on device. This does not affect factory function calls which are called with an explicit device 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.

Warning

This function imposes a slight performance cost on every Python call to the torch API (not just factory functions). If this is causing problems for you, please comment on https://github.com/pytorch/pytorch/issues/92701

Note

This doesn’t affect functions that create tensors that share the same memory as the input, like: torch.from_numpy() and torch.frombuffer()

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)