ptrblck:

torch.matmul(A.permute(0, 1, 3, 2), B)


As now I am getting this error:
TypeError: matmul(): argument ‘input’ (position 1) must be Tensor, not Softmax
At line:
 attention = torch.matmul(attention, V.permute(0, 1, 3, 2))
This is for second matmul() method