pytorch指定gpu

在 PyTorch 中，可以使用以下代码将模型和张量移动到特定的 GPU 设备上：

import torch

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    device = torch.device("cuda")  # 默认选择第一个可用的 GPU
    # 或者使用指定的 GPU 设备
    # device = torch.device("cuda:0")  # 选择特定的 GPU 设备
else:
    device = torch.device("cpu")

# 将模型加载到指定的设备上
model.to(device)

# 将张量移动到指定的设备上
tensor = tensor.to(device)

如果你有多个 GPU，可以使用 `cuda:0`, `cuda:1`, 等来选择特定的 GPU 设备。

请确保安装了与你的 PyTorch 版本相匹配的 CUDA 和 cuDNN 版本，并且已正确配置 CUDA 路径。