PyTorch torch.device Function
PyTorch torch Reference Manual
torch.deviceis a function in PyTorch used to create device objects, used to specify the device on which tensors are stored.
Function Definition
torch.device(device)
Usage Examples
Examples
import torch
# Create CPU device
cpu_device = torch.device('cpu')
print("CPU device:", cpu_device)
# Create CUDA device
if torch.cuda.is_available():
cuda_device = torch.device('cuda:0')
print("CUDA device:", cuda_device)
# Create tensor on GPU
x = torch.zeros(3, 4, device=cuda_device)
print("Tensor device:", x.device)
# Create CPU device
cpu_device = torch.device('cpu')
print("CPU device:", cpu_device)
# Create CUDA device
if torch.cuda.is_available():
cuda_device = torch.device('cuda:0')
print("CUDA device:", cuda_device)
# Create tensor on GPU
x = torch.zeros(3, 4, device=cuda_device)
print("Tensor device:", x.device)
Other Extensions