I am currently working a project like this. Although I cannot provide you the actual code, I can provide you the blog post and reference that actually do this.

• https://jacobgil.github.io/deeplearning/vision-transformer-explainability
• https://www.kaggle.com/code/piantic/vision-transformer-vit-visualize-attention-map

Both post provided code examples and the first one goes more into the theory behind this. I hope this helps.