How to save and load nn.Parameter() in pytorch, such that one can continue optimisation over it?
I know how to store and load nn.Model, but can not find how to make a checkpoint for nn.Parameter. I tried this version, but the optimizer is not changing the nn.Parameter value after restoring.
from torch import nn as nn
import torch
from torch.optim import Adam

alpha = torch.ones(10)
lr = 0.001
alpha = nn.Parameter(alpha)
print(alpha)
alpha_optimizer = Adam([alpha], lr=lr)

for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)
path = &quot;./test.pt&quot;
state = dict(alpha_optimizer=alpha_optimizer.state_dict(),
               alpha=alpha)
torch.save(state, path)
checkpoint = torch.load(path)
alpha = checkpoint[&quot;alpha&quot;]
alpha_optimizer.load_state_dict(checkpoint[&quot;alpha_optimizer&quot;])
for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)


The problem is that the optimizer is still with a reference to the old alpha (check id(alpha) vs. id(alpha_optimizer.param_groups[0][&quot;params&quot;][0]) before the last for loop), while a new one is set when you load it from the checkpoint in alpha = checkpoint[&quot;alpha&quot;].
You need to update the params of the optimizer before loading its state:
# ....
torch.save(state, path)
checkpoint = torch.load(path)

# here's where the reference of alpha changes, and the source of the problem
alpha = checkpoint[&quot;alpha&quot;]

# reset optim
alpha_optimizer = Adam([alpha], lr=lr)
alpha_optimizer.load_state_dict(checkpoint[&quot;alpha_optimizer&quot;])

for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)

