PyTorch torch.set_default_device Function
Pytorch torch Reference Manual
torch.set_default_deviceIt is a function in PyTorch used to set the default device.
Function Definition
torch.set_default_device(device)
Usage Example
Example
import torch
# Set the default device
if torch.cuda.is_available():
torch.set_default_device('cuda')
x = torch.zeros(3, 4)
print("Default device:", x.device)
torch.set_default_device('cpu')
else:
print("CUDA is not available")
# Set the default device
if torch.cuda.is_available():
torch.set_default_device('cuda')
x = torch.zeros(3, 4)
print("Default device:", x.device)
torch.set_default_device('cpu')
else:
print("CUDA is not available")
Other Extensions