Visualizing co-attention

[rakeshshrestha31]

I was wondering if it is possible to visualize the co-attention similar to figure 3 of the paper (also shown below). Is there a specific tensor we can use or does it require further processing?