Visualize attention maps for ViT

[apoorvasafai]

Hi

I am using ViT for a 2D image classification task. How do I visualize the attention maps after changing save_attn=True while training the ViT model.

Thanks

[KumoLiu]

Hi @apoorvasafai, you can get it by using:
net.blocks[0].attn.att_mat
Hope it helps, thanks.