这道题在恢复 checkpoint 后的 optimizer.step() 失败，堆栈落在 SGD 动量更新，而不是模型前向：
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
为什么前向通过，step 还会失败
原帖把模型、输入与损失函数移到 GPU 后，已经走过了前向和反向；异常出现在 buf.mul_(momentum).add_(d_p, ...)。这里参与运算的是动量缓冲区和当前梯度，不只是图片与权重。检查输入设备固然有用，却覆盖不到恢复后的优化器状态。
原帖没有贴出完整加载顺序，所以“状态留在 CPU”是要用下面打印结果确认的分支。
因此即使输入、标签、模型都已上 GPU，优化器保存的动量仍值得单独检查。将下列代码放在第一次 optimizer.step() 前；optimizer 就是恢复状态后的优化器。
import torch
print("josh050608")
for group in optimizer.param_groups:
for p in group["params"]:
print("param/grad:", p.device, None if p.grad is None else p.grad.device)
for key, value in optimizer.state.get(p, {
}).items():
if isinstance(value, torch.Tensor):
print("state:", key, value.device, value.dtype)
若某参数在 CUDA、其 momentum_buffer 却在 CPU，就定位到了这次加法的设备冲突。优先调整恢复顺序：把模型放到目标设备，再以该模型参数创建优化器，最后加载优化器状态。checkpoint_path 是自己的可信 checkpoint 路径，Model 和优化器超参数需保持原训练定义。
device = torch.device("cuda:0")
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
model = Model().to(device)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
旧版 PyTorch 没有 weights_only 参数；按该版本签名调用。原帖 checkpoint 还保存了 args 等对象，若受限加载不接受它们，应在可信来源前提下按对应版本处理，而不是对未知文件关闭检查。
对 SGD，必要时仅把确认错放的 momentum_buffer 迁到对应 p.device。不要把所有优化器的状态一律搬动，一些优化器允许 CPU 上的标量计数。再做一次前向、反向和 step，重复打印确认动量与参数一致，并确认参数确实发生更新。
只修已经确认错放的 SGD 动量
若检查证实只有 SGD 的动量设备不对，在 load_state_dict 后、下一次 step 前可以显式修正：
# optimizer 是前文恢复好的 SGD；不适用于任意优化器的所有状态。
for group in optimizer.param_groups:
for p in group["params"]:
state = optimizer.state.get(p, {
})
buf = state.get("momentum_buffer")
if buf is not None and buf.device != p.device:
state["momentum_buffer"] = buf.to(device=p.device)
同时确认优化器实际持有的就是当前模型参数。若先建优化器，随后又把 model 重新赋值为另一个实例，优化器可能仍在更新旧对象；单纯迁移动量不能修复这种引用错误：
model_ids = {
id(p) for p in model.parameters()}
optimizer_ids = {
id(p) for g in optimizer.param_groups for p in g["params"]}
print("optimizer belongs to model:", optimizer_ids <= model_ids)
下面的小例子完整演示保存、重建、恢复与一次更新。使用内存文件和 CPU，不读写真实 checkpoint；它验证恢复流程，不复现 GPU 设备冲突：
import io
import torch
from torch import nn
torch.manual_seed(3)
model = nn.Linear(3, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
x, y = torch.randn(4, 3), torch.randn(4, 1)
optimizer.zero_grad()
nn.functional.mse_loss(model(x), y).backward()
optimizer.step()
buffer = io.BytesIO()
torch.save({
"model": model.state_dict(), "optim": optimizer.state_dict()}, buffer)
buffer.seek(0)
checkpoint = torch.load(buffer, map_location="cpu", weights_only=True)
restored = nn.Linear(3, 1).to("cpu")
restored.load_state_dict(checkpoint["model"])
opt = torch.optim.SGD(restored.parameters(), lr=0.1, momentum=0.9)
opt.load_state_dict(checkpoint["optim"])
before = restored.weight.detach().clone()
opt.zero_grad()
loss = nn.functional.mse_loss(restored(x), y)
loss.backward()
opt.step()
assert not torch.equal(before, restored.weight)
for p, state in opt.state.items():
assert state["momentum_buffer"].device == p.device
print("restored step:", float(loss.detach()))
回到 GPU 环境后，用同一小批真实数据检查梯度、动量、参数三者的设备，再确认权重变化。不要用“删除优化器状态就能跑”作为等价的断点续训验证，那会丢掉原来的动量历史。
参考：原始求助帖；技术依据：PyTorch 优化器（2.5）、PyTorch 保存与加载教程、torch.load（2.5）。
