pytorch 如何使用cuda

首先,在最上面放进去

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

其次,网络进入device

net.to(device)

最后,数据进入device

inputs, labels = inputs.to(device), labels.to(device)

以上在运行程序的时候都没有出现问题,但在输出的时候报错了

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

解决方式:

images, labels = data
# 在输出的时候出现的问题,我应该把image也进到cuda里面,然后出来的是inputs,这时候再把
# inputs 输入到net里面,总是,要用cuda就全都用cuda,否则都不用
inputs, labels = images.to(device), labels.to(device)
outputs = net(inputs)

也就是说,在验证使用的时候,还是要把image进入到device里面,然后再进入net里面,总之,所有的要么都进GPU,要么都不进GPU

猜你喜欢

转载自blog.csdn.net/huatianxue/article/details/108273309