现在位置: 首页 > PyTorch 教程 > 正文

PyTorch torch.device 函数


Pytorch torch 参考手册 Pytorch torch 参考手册

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)

Pytorch torch 参考手册 Pytorch torch 参考手册