Solved by AladdinPerzon in post #2 
              
              
                I think you can use einsum see below for a small example. It felt a bit confusing doing dot product over multiple dimensions but it seems to make sense to me after a few examples. 
N, M, D = 2, 3, 4
x = torch.ones((N,M,D))
y = torch.ones((N,M,D))
torch.einsum('ijk,ijk->ij', x, y).unsqueeze(2) # will…