Thanks a lot @111414 .
I also found that swapping the position between mat1 and mat2 also works:
batch_size = 256
mat1 = torch.randn(25,batch_size)
mat2 = torch.randn(36,36,25)
torch.matmul(mat2, mat1).shape