Solved by albanD in post #2 
              
              
                Hi, 
I think the best way to do this is to rewrite (assuming 2D matrices) matmul(m1, m2).diag() to (m1 * m2.t()).sum(dim=1). 
You can handle the extra batch dimension if you have one by just changing the the dimension you sum over.