Fix Grad-CAM heatmap generation

[NateBioTech]

Added attention weights handling in src/stamp/modeling/vision_transformer.py

[s1787956]

I am confused. Where are the attention weights used in the gradcam function?
In general, there are two different approaches for Transformer heatmaps:

1. gradient back to input (our standard nowadays)

2. attention activation (attention weights) visualization

3. has the disadvantage that parts of the model are kinda neglected (MLP parts). And there are multiple heads and layers that have individual activations).

So I am unsure if we are not mixing these things with this approach. Also returning the weights will affect the performance of the ViT.