背景介绍
我们在使用Pytorch训练时，模型和数据有可能加载在不同的设备上（gpu和cpu），在算梯度或者loss的时候，报错信息类似如下：

RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' in call to _thnn_l1_loss_forward

这种报错通常有几种情况：

  1. 数据在cpu上，模型在gpu上；

  2. 数据在gpu上，模型在cpu上；

  3. 指定层不支持gpu或cpu版本，但数据和模型都在gpu或cpu上。

解决方案
第1步：将模型和数据都移到同一设备上。

假设模型和数据定义如下：

model = Model()

input = dataloader()

output = model(input)

移动到gpu上：

# solution: 0

device = 'gpu'

model = model.to(device)

data = data.to(device)


# solution: 1

model = model.cuda()

data = data.cuda()

移动到cpu上：

# solution: 0

device = 'cpu'

model = model.to(device)

data = data.to(device)


# solution: 1

model = model.cpu()

data = data.cpu()

第二步：打印模型model和数据data在gpu上还是cpu上。

  1. 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。 
print('Is model on gpu: ', next(model.parameters()).is_cuda)
输出若是True，则model在gpu上；若是False，则model在cpu上。
  1. 输出数据data的device字段。 
print('data device: ', data.device)
输出gpu则在gpu上，输出cpu则在cpu上。