Just want to mention that one should be cautious when using this line
grad_dict = {k:v.grad for k, v in zip(net.state_dict(), net.parameters())}

as net.state_dict() 's keys is not guranteed to be in the same order as in net.parameters(), as I tried this in my SSD model. So I would sugget to use:
grad_dict = {k:v.grad for k,v in net.named_parameters()}

(Maybe in that time(2 years ago), we don’t have this issue or named_parameters is not available, I’m not sure, but I think it’s better to use named_parameters nowadays)