Issue loading mechanic optimizer state

[yitzhaklevi]

The following simple script reproduces the issue:

model = torch.nn.Linear(1, 1)
optimizer = mechanize(torch.optim.AdamW)(model.parameters(), lr=1e-5)
x = torch.ones([5, 1])
out = torch.sum(model(x))
out.backward()
optimizer.step()
print('done first step')

new_optimizer = mechanize(torch.optim.AdamW)(model.parameters(), lr=1e-5)
new_optimizer.load_state_dict(optimizer.state_dict())
out = torch.sum(model(x))
out.backward()
new_optimizer.step()
print('done new steps using new optimizer loaded')

[yitzhaklevi]

The issue is due to the fact that state_dict['state']['_mechanic'] has tensor pointers as keys, those does not have the same addresses when re-initializing new Mechanic optimizer

(other optimizers e.g AdamW has the indexes as keys for the state)

[acutkosky]

Thanks for bringing this up! I'll take a look and make an update (or if you want to do so, please feel free to submit a PR).

[yitzhaklevi]

Welcome, I actually fixed that (will submit a PR next week, ) but it seems that on Windows the issue does not reproduce. (tried on my local machine and it worked)

[ogencoglu]

Any update on this?