One thing is that this:
for param_group in optimizer.state_dict()['param_groups']:

should be replaced with that:
for param_group in optimizer.param_groups:

I know that the first version appears in the ImageNet example, but it no longer works as expected. I can’t see anything wront at a glance, but I’ll try to look more carefully into it sometime.