I have this problem too, is it reasonable to reset optimizer.state=collections.defaultdict(dict) like:
import torch
import torch.nn as nn
from torch import optim
import collections

m = nn.Linear(3, 1)
opt = optim.Adam(m.parameters(), lr=1e-3)
out = m(torch.rand(3))
out.backward()
opt.step()
print(opt.state)

opt.state = collections.defaultdict(dict) # Reset state
print(opt.state)