How to record each layer's gradients in each epoch?

#1

How to record each layer’s gradients in each epoch?

#2

You can use register_backward_hook and gather the grads in some shared dict