Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

在Pytorch上做QAT训练时，最开始使用的CPU，训练正常，后来尝试在GPU上训练，于是改动代码如下：

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

在训练时出现题目中错误：

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

看意思应该是有些tensor在cuda上，有些tensor在cpu上，需要将两者统一起来。

检查了一下代码，发现训练数据都已经load到了设置的device上：

# Copy data into device
noise = noise.to(device=device, dtype=torch.float32)
target = target.to(device=device, dtype=torch.float32)

checkpoints也映射到了对应的device上 ：

checkpoints = torch.load(model_path, map_location=device)

那么就要检查一下模型是在那个device上了，加一行打印：

print(next(model.parameters()).device)

发现打印出来是CPU，说明模型还在CPU上，没有load到cuda。添加语句加载到cuda上：

model.to(device)

问题解决。