'''
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, …?