Finding attention scores for self-attention visualisation

[nigelnnk]

Hi there! In the MAT paper, there were visualisations of the self-attention weights that were produced for one molecule. I am a bit unsure of where these self-attention weights can be located in the transformers.py file. Do you have any advice? Thanks!