import torch
import torch.nn as nn

X = torch.rand(2, 10)
y = torch.tensor([[0,1], [1,0]], dtype=torch.float32)

m = nn.Linear(10, 2)
opt = torch.optim.Adam(m.parameters())
state = opt.state_dict()
crit = nn.BCEWithLogitsLoss()

for i in range(5):
    if i == 2:
        # load
        opt.load_state_dict(state)
    opt.zero_grad()
    out = m(X)
    loss = crit(y, out)
    loss.backward()
    opt.step()
    # save
    state = opt.state_dict()

Update to less dummy example. Also works fine.