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