PyTorch torch.device 函数
torch.device 是 PyTorch 中用于创建设备对象的函数,用于指定张量存储的设备。
函数定义
torch.device(device)
使用示例
实例
import torch
# 创建 CPU 设备
cpu_device = torch.device('cpu')
print("CPU 设备:", cpu_device)
# 创建 CUDA 设备
if torch.cuda.is_available():
cuda_device = torch.device('cuda:0')
print("CUDA 设备:", cuda_device)
# 在 GPU 上创建张量
x = torch.zeros(3, 4, device=cuda_device)
print("张量设备:", x.device)
# 创建 CPU 设备
cpu_device = torch.device('cpu')
print("CPU 设备:", cpu_device)
# 创建 CUDA 设备
if torch.cuda.is_available():
cuda_device = torch.device('cuda:0')
print("CUDA 设备:", cuda_device)
# 在 GPU 上创建张量
x = torch.zeros(3, 4, device=cuda_device)
print("张量设备:", x.device)

Pytorch torch 参考手册