Solved by tom in post #2 
              
              
                Either use einsum with bij,bij->bi or use unsqueeze(-2) and unsqueeze(-1) on the inputs of matmul and squeeze the result. (matmul reduces the second to last and last axes of the first and second input, respectively) or you could multiply with * and then sum at the expense of some time and memory.