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