pytorch无论是在模型的训练和测试阶段,都需要将模型,数据,标签这三者放在同一设备上,如果使用cpu,则均需在cpu,如果使用gpu,则均需在gpu中,我们通常使用以下代码进行对齐:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
data = Variable(data.to(device))
labels = Variable(labels.to(device))