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