how can we get attention weights for visualization

Hi. Please i am very interested in your code and i would like to know how to extract weights for attention visualization. If possible how to do visualization. Your code is great and thanks a lot.