RuntimeError: The size of tensor a (128) must match the size of tensor b (512) at non-singleton dimension 2

[OmarMohammed88]

I followed your instructions and replaced class MultiheadAttenion with LinearMultiheadAttention and kept the seq_len 512 and proj_k 128, my configurations were as follow ( hidden dim 512, max_text length 512) .
After more Debuging i found that
attn_mask is torch.Size([1, 512, 512])
attn_output_weights torch.Size([64, 512, 128]) .