Visualizing the attention weights

[ahof1704]

Hi,

First of all, thank you so much for the great work you guys have done in your scripts.
I would like to visualize the attention weights obtained after training. I tried to use this visualization tool: https://github.com/jessevig/bertviz#attention-head-view

However, it seems to only work for the pre-trained BERT model, not the fine-tuned one we make ourselves via your script.

What would you recommend for visualization?

Thank you!