BART: get activation maps for encoder and decoder

Hello, I am using BART for abstractive text summarization in Italian language and I would like to extract the activation (or feature) maps for diagnosis purpose.

I would also like to store and check them in TensorBoard, or other monitoring tools, during training. I already checked the TrainerCallback on the docs and I think that is the one I need to subclass in order to extract the activation maps.

Any help / redirect / etc. to perform this task will be appreciated. Thanks in advance!