I am using attention weighted pooling strategy to generate cross-attentions across N transformer heads, this allows us to visualize which parts of an output wer

I am using attention weighted pooling strategy to generate cross-attentions across N transformer heads, this allows us to visualize which parts of an output were caused by which part of the input -- this is used to debug summarization. Only works for encoder/decoder networks - i.e. ones where you have cross attentions