Is there a better way to multiply & sum two Pytorch tensors along the first dimension?
I have two Pytorch tensors, a &amp; b, of shape (S, M) and (S, M, H) respectively. M is my batch dimension. I want to multiply &amp; sum the two tensors such that the output is of shape (M, H). That is, I want to compute the sum over s of a[s] * b[s].
For example, for S=2, M=2, H=3:
&gt;&gt;&gt; import torch
&gt;&gt;&gt; S, M, H = 2, 2, 3
&gt;&gt;&gt; a = torch.arange(S*M).view((S,M))
tensor([[0, 1],
        [2, 3]])
&gt;&gt;&gt; b = torch.arange(S*M*H).view((S,M,H))
tensor([[[ 0,  1,  2],
         [ 3,  4,  5]],

        [[ 6,  7,  8],
         [ 9, 10, 11]]])

'''
DESIRED OUTPUT:
= [[0*[0, 1, 2] + 2*[6, 7, 8]], 
   [1*[3, 4, 5] + 3*[9, 10, 11]]]

= [[12, 14, 16],
   [30, 34, 38]]

note: shape is (2, 3) = (M, H)
'''

I've found one way that sort of works, using torch.tensordot:
&gt;&gt;&gt; output = torch.tensordot(a, b, ([0], [0]))
tensor([[[12, 14, 16],
         [18, 20, 22]],

        [[18, 22, 26],
         [30, 34, 38]]])
&gt;&gt;&gt; output.shape
torch.Size([2, 2, 3]) # always (M, M, H)
&gt;&gt;&gt; output = output[torch.arange(M), torch.arange(M), :]
tensor([[12, 14, 16],
        [30, 34, 38]])

But as you can see, it makes a lot of unnecessary computations and I have to slice the ones that are relevant for me.
Is there a better way to do this that doesn't involve the unnecessary computations?

This should work:
(torch.unsqueeze(a, 2)*b).sum(axis=0)
&gt;&gt;&gt; tensor([[12, 14, 16],
            [30, 34, 38]])

