Torch.tensordot - Dot product in Tensor along 3rd dim

1

I am two tensors A, B, each of shape n x n x m and I would like to get a tensor with dimensions n x n where each entry is the dot product of long the third dimension, i.e. $$C_{ij}=\sum_k A_{ijk}B_{ijk}$$.

Example:

A=B=[[[ 0.,  1.,  2.,  3.,  4.], [ 5.,  6.,  7.,  8.,  9.]],
        [[10., 11., 12., 13., 14.], [15., 16., 17., 18., 19.]]]

I would like to get

[[  30.,  255.],
[ 730., 1455.]]

Code:

a = torch.arange(20.).reshape(2, 2, 5)
b = torch.arange(20.).reshape(2, 2, 5)

to

for i in range(2):
    for j in range(2):
        d[i,j] = a[i,j].dot(b[i,j])

using a single operation

2

Overlooked the most obvious way to solve this: c = (a*b).sum(2)