Can't find momentum buffer for Linear layer

Hi,
I have a problem with my code. The error first emerged after upgrading to 1.7.0

I used

for name, param in model.named_parameters():
mom_param = optimizer.state[param][‘momentum_buffer’]

before to get the momentum buffer. As Iread this is no longer working in 1.7.0

For all except one layer (Linear) I can use:

k = 0
for name, param in model.named_parameters():
mom_param = optimizer.state_dict()[k][param][‘momentum_buffer’]

What can I do to get the momentum buffer of my Linear layer?

I seachred hot to make the code inside one of this fancy code snips. But I haven’t found out how to do it.

Thanks for helping.
Jessi

The first code snippet seems still to work in 1.8.0.dev20201116:

model = models.resnet18()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
model(torch.randn(1, 3, 224, 224)).mean().backward()
optimizer.step()

for name, param in model.named_parameters():
    mom_param = optimizer.state[param]['momentum_buffer']
    print(mom_param)

Could you share a code snippet, where this would be failing?