optimizer.step() Expected all tensors to be on the same device, but found at least two device

这个错误通常是由于模型参数和优化器的设备不匹配引起的。为了解决这个问题，你可以采取以下几种方法：

1. 确保模型和优化器在同一设备上：

   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
   optimizer = optimizer_class(model.parameters(), lr=learning_rate)

2. 检查模型参数的设备：

   print(model.parameters())

3. 检查输入数据的设备：

   inputs = inputs.to(device)
   labels = labels.to(device)

4. 如果你使用了多个GPU进行训练，确保模型和优化器都在同一个GPU上：

   model.to("cuda:0")
   optimizer = optimizer_class(model.parameters(), lr=learning_rate)

5. 如果你的模型和优化器已经在同一设备上，但仍然出现错误，可能是因为你的模型参数在不同的设备上。你可以使用以下代码将所有模型参数移动到同一设备上：

   model.to("cuda:0")
   model = nn.DataParallel(model)