Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_mm)

这个错误通常是因为你的输入数据在不同的设备上。在PyTorch中，每个张量都必须是相同的设备类型（CPU或GPU）。因此，如果你的模型在GPU上运行，而输入数据在CPU上，就会出现这个错误。

解决方法是将所有输入数据都移到同一个设备上。你可以使用`.to()`方法将张量移动到特定的设备上。例如，使用`tensor.to('cuda')`将张量移动到GPU上，使用`tensor.to('cpu')`将张量移动到CPU上。