Hello good people .
Let’s say I have loaded the Vit with pretrained = True. And I want to extract the output of softmax of the attention heads from this model. How can I achieve that? should I create a new model and inherit from this model and then change the forward method. Or should I use the modules for the Vit and load each layer to the new model?