I think we will have to modify F.multi_head_attention_forward for that,
in the end it has,
if need_weights:
        # average attention weights over heads
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        return attn_output, attn_output_weights.sum(dim=1) / num_heads
else:
        return attn_output, None

if we want full attention weights without averaging over heads, then we will have to change it to,
if need_weights:
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        return attn_output, attn_output_weights