'''
Let's say we have a 2D tensor A and a 3D tensor B.
Let A.size() equals (2, 3) and B.size() equals (2, 3, 3)
I expect the result C to be a 2D tensor with C.size() equals (2, 3).
'''
C[0] = torch.matmul(B[0], A[0])
C[1] = torch.matmul(B[1], A[1])

Can I vectorize the above operation instead of looping over 0, 1, …?

'''
Seems I can leverage torch.bmm to get it done !!
torch.bmm can take two tensors of (b, n, m) and (b, m, p)
to give a tensor of (b, n, p)
So, let me put A in the proper shape to leverage torch.bmm
'''
A = torch.reshape(2, 3, 1)
C = torch.bmm(B, A).reshape(2, 3)