similar behaviour as output_attentions for a custom finetuned model

I am fine tuning a BERT classifier with my own data and save the state_dict (following this tutotrial). Now for evaluating the model, I would like to use bertViz, and so one thing I need is for an option like huggingface’s “output_attentions=True” for a custom model class but I’m not sure how to do that.