Thanks for the clarification.

A stateless optimizer might work, but if it’s using internal buffers, such as e.g. Adam, you would get a device mismatch error. I would generally recommend to recreate the optimizer and load its state_dict. An annoyance is that you would need to push the internal states to the new device. Here is a code snippet adapted from this issue 1 (a to() method for optimizers might be useful here  ):

model = nn.Linear(1, 1).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Single update
out = model(torch.randn(1, 1).cuda())
out.mean().backward()
w0 = model.weight.clone()
optimizer.step()
w1 = model.weight.clone()
print((w1 - w0).abs().max())
optimizer.zero_grad()

# Move to CPU and verify update
sd = optimizer.state_dict()
model.cpu()
optimizer_cpu = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.load_state_dict(sd)
for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.cpu()
            
out = model(torch.randn(1, 1))
out.mean().backward()
w0 = model.weight.clone()
optimizer.step()
w1 = model.weight.clone()
print((w1 - w0).abs().max())


Yes, the deepcopy seems to be fine here. The load_state_dict method would call param.copy_ as seen here 2, so it shouldn’t be necessary to create the deepcopy of the state_dict, but it might not hurt.