error in self attention visualisation

I beleive there is a error for idexing the self attetion:

for idx_o, ax in zip(idxs, axs):
    idx = (idx_o[0] // fact, idx_o[1] // fact)
    ax.imshow(sattn[..., idx[0], idx[1]], cmap='cividis', interpolation='nearest')
    ax.axis('off')
    ax.set_title(f'self-attention{idx_o}')

should be sattn[idx[0], idx[1], ...]as the last two dim sums to 1, so the first two are locations and are the dim to indexing, with this the attention map is

[fmassa]

@songhwanjun good catch. I had fixed this for the class AttentionVisualizer, but looks like I missed it in the example cell just above. I’ll send a PR fixing it

[fmassa]

Hi,

This is a fair point and is discussed in the bottom of the notebook as well

In the end, it’s a matter of deciding if you want the attention of all pixels at location x, or the contribution of location x to all the attention maps.

I would say both information are useful to be considered, and as such I added both in the widget visualization.