Pytorch中实现CPU和GPU之间的切换

 

案例：

# 判断是否可以使用GPU
if torch.cuda.is_available():
    d = torch.device('cuda')  
    a = torch.tensor([[1, 2], [3, 4]])
    b = torch.ones_like(a, device=d)  # b在GPU上操作
    c = a.to(device=d)  # c在GPU上操作
    d = torch.add(a, b)
    print(d)
    print(d.to('cpu', torch.float64))