Efficient method to compute the row-wise dot product of two square matrices of the same size in PyTorch
Supposing I have two square matrices A, B of the same size
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[1, 1], [1, 1]])

And I want a resulting tensor that consists of the row-wise dot product, say
tensor([3, 7])  # i.e. (1*1 + 2*1, 3*1 + 4*1) 

What is an efficient means of achieving this in PyTorch?

As you said you can use torch.bmm but you first need to broadcast your inputs:
&gt;&gt;&gt; torch.bmm(A[..., None, :], B[..., None])
tensor([[[3]],

        [[7]]])

Alternatively you can use torch.einsum:
&gt;&gt;&gt; torch.einsum('ij,ij-&gt;i', A, B)
tensor([3, 7])

