What is the proper way to use optim.ASGD?

1

When using the optimizer ASGD, the optimizer stores the average parameter values in optimizer.state_dict()['state'][...]['ax']. Is there a natural way to perform evaluation with this model ? At the moment I am using

tmp = {}
for p in model.parameters():
    tmp[p] = p.data
    p.data = optimizer.state[p]['ax']

evaluate(model)

for p in model.parameters():
    p.data = tmp[p]

But I feel weird cloning my model at every evaluation and I am wondering wether this could be done with some pretty context manager.