Attention maps in Different query position

[kbiyani33]

Hello,
Is it possible to view Attention maps at different query positions ?
If yes, how can we do it?

[xvjiarui]

You may follow this simple script to draw the heat map.
For NonLocal, you need to modify the code.