Hi, Based on your error, you would have to transpose one matrix for them to be multiple, as the second dimension of the first matrix and the first dimension of the second matrix should be the same.
Therefore you should multiply matrices of shape 128 x 32 and 32 x 128 to get a resultant shape of 128 x 128
x_t = a_t @ x_t

because a_t would be of shape batch x hidden_size and x_t is of shape batch x input_size
Please check your input shapes once.  I believe the error comes from the line
On a side note, instead of using the @ operator, I believe using torch.bmm would be much faster.