PyTorch Optimizer报RuntimeError：张量设备不匹配问题求助

各位大佬好！我最近在写PyTorch训练代码的时候卡壳了，一直报设备不匹配的错误，折腾了好久都没搞定，想请大家帮忙看看。

先贴一下我的训练代码片段：

# create a Pytorch optimizer 
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):
    if iter % eval_iters == 0:
        losses = estimate_loss()
        print(f"step: {iter}, loss {losses}")
        
    # sample a batch of data
    xb, yb = get_batch("train")
    
    #evaluate the loss
    logits, loss = model.forward(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    
print(loss.item())

运行这段代码的时候，直接抛出了这个错误：

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

我已经试过把模型移到指定设备上（model = model.to(device)），也试过给输入张量加.cuda()，但都没解决问题，现在实在不知道该怎么排查了，有没有大佬能指点一下？

备注：内容来源于stack exchange，提问作者user29981993