How do I obtain the attention, query and key of fine tuned model output?

Hello all,

I am trying to visualise attention throughout the fine-tuned BioMegatron Model with token classification(ner) and text classification(re) downstream tasks using bertviz.

How can I get the attention, query, and key sequences? It seems that the NeMo code cannot be edited on Colab.

Can anyone assist?

Thanks,