Hello,
I am a bit confused on why for batch size >1 in the MultiHeadAttention class, attention is calculated by passing the input as <batch_size*num_heads, seq_length, embed_dim> and later reshaped to < batch_size, num_heads, seq_length, embed_dim> to get average attention across heads. Shouldn’t the input be reshaped first to split by batch and then calculate attention?
Thank you!