Want get Attention map

[Hwanman]

Hi. The code runs fine. But I am learning Vit and want to check the attention map. Where should I modify it? First of all, the current code doesn't seem to output attention weights.