Multiplying every row of a 2D tensor with every matrix in a 3D tensor

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

2 
'''
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)