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?