最近我运行了一个简单的代码出现了下面的错误
Expected object of backend CUDA but got backend CPU for argument #4 'mat1'
过程是这样的,我建立了一个模型model = LeNet()
,然后判断是否有cuda
,如果存在cuda
,然后将模型迁移到cuda
上
if torch.cuda.is_available():
model.cuda()
在模型训练中也同样的操作,在dataloader中提取出来的数据,然后放入cuda
中,但是后面出现了上面的错误。
for (data, label) in tqdm(dataloader):
data = data.view(data.size(0), -1)
data = torch.autograd.Variable(data, requires_grad = True)
if torch.cuda.is_available():
data.cuda()
原因在于模型迁移到cuda
的命令model.cuda()
可以将模型复制到GPU上,且改变变量model
作为GPU内存的变量,后面引用model
的时候用的都是GPU
上的数据,而data.cuda()
却正好相反,只有返回值才能在GPU上,所以正确的做法是data = data.cuda()
造成上述的差异的原因是torch.nn.Module.cuda()
和torch.Tensor.cuda()
的问题,同样会出现问题的是torch.nn.Module.to()
和torch.Tensor.to()