Pytorch使用GPU计算

方法一

可以使用.cuda()将数据转换到GPU中,用.cpu()再换到CPU中

x = torch.zeros(3,3,dtype=torch.double)
y = torch.ones(3,3,dtype=torch.double)
if torch.cuda.is_available():
    x = x.cuda()
    y = y.cuda()
    z = x+y
    print(z)

#将x放回CPU中
x = x.cpu()

方法二

还可以使用.to()方法进行CPU和GPU的转换

x = torch.zeros(3,3,dtype=torch.double)
y = torch.ones(3,3,dtype=torch.double)
device = torch.device("cuda")
print(device)

x = x.to(device)
y = y.to(device)
z = x+y

另外，在Linux上实时刷新显卡使用情况用指令 watch -n 1 -d nvidia-smi