是用pycharm debug如下代码即可知道。
torch.cuda.set_device(0)会修改torch.device('cuda')的值,即默认设备。
仅允许同一设备内部的tensor进行计算;

import torch

# 检查 GPU 是否可用
if torch.cuda.is_available():
    # 使用默认 CUDA 设备
    torch.cuda.set_device(0)
    x = torch.tensor([1.0, 2, 3], device=torch.device('cuda'))
    print("默认 CUDA 设备上的 Tensor:", x)

    # 设置当前设备为第一个 GPU
    torch.cuda.set_device(1)
    c = torch.tensor([1.0, 2, 3], device=torch.device('cuda'))
    b=torch.cuda.current_device()
    x2=x+x
    x3 = x2.to(1)
    print("当前设备:", torch.cuda.current_device())
else:
    print("没有可用的 GPU")

最后修改:2023 年 06 月 27 日
如果觉得我的文章对你有用,请随意赞赏