How to extract gradient during training in pytorch with Trainer module?

Hello,

I am quite familiar overall with the Trainer module and the models. Yet, it is not perfectly clear to me how to customize it to get gradient metrics like the norm by layer. What would be the best way?

Thanks in advance for your help!