Hi,
Keep in mind that state_dict() does not copy your weights. So after you update your model to contain W2, your state dict also contains W2.
You might want to do it by hand with saved_params = [p.clone() for p in model.parameters()] and load it with
with torch.no_grad():
  for p, old_val in zip(model.parameters(), saved_params):
    p.copy_(old_val)