Hadamard-Multiply every image in batch with every other

So suppose I have two tensors of the following shape: (batch_size, n, m), name them A and B.

What I want is to build a tensor of shape (batch_size**2, n, m) such that every row of A is entry-wise multiplicated with every row of B, i.e., the first batch_size rows should consistent of the first row of A times every row in B, the next batch_size rows should consequently be the second row of A times every row in B, etc.

How would one do this in Pytorch? Can this be solved by simple Broadcasting?
Thanks a lot!

Found a solution myself using torch.einsum:

torch.einsum('aij,bij->abij', A, B)