multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
attn_output, attn_output_weights = multihead_attn(query, key, value)
'attn_output_weights` only contains the softmax-normalized attention scores. Is this true?
Or it is the final output after multiplication with the value?