Solved by tom in post #3 
              
              
                There are several ways to get this easily, for example 
b = torch.einsum('bii->b', a)

or 
b2 = torch.diagonal(a, dim1=-2, dim2=-1).sum(-1) 

Best regards 
Thomas