PyTorch 中将 Tensor 移至 CPU 或 GPU 设备

PyTorch 中将 Tensor 移至 CPU 或 GPU 设备

此段代码展示了如何将一个 PyTorch Tensor 对象移至 CPU 或 GPU 设备：

import torch

# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tensor_on_device = some_tensor.to(device)
tensor_on_gpu = tensor_on_cpu.to('cuda')

定义 'some_tensor' 内容：

'some_tensor' 可以是任何一个 torch.Tensor 对象。可以是一个随机初始化的张量，也可以是从数据中加载的张量。下面是一些示例：

1. 随机初始化一个张量：

some_tensor = torch.randn((3, 3))

2. 从数据中加载一个张量：

import numpy as np

data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
some_tensor = torch.from_numpy(data)

3. 使用 torch.zeros() 函数创建一个全零张量：

some_tensor = torch.zeros((2, 2))

4. 使用 torch.ones() 函数创建一个全一张量：

some_tensor = torch.ones((2, 2))

无论你选择哪种方式，都可以将 'some_tensor' 移动到所需的设备上（CPU 或 GPU）。