Hey! m_alpha_beta contains input, which is basically just a b x n x n parameter tensor, where b is the batch size. m_beta_alpha contains input.permute(0, 2, 1), which is basically the transpose of m_alpha_beta.
I tested it with the input above, which is
input = torch.Tensor([[[7,  8, 7, 20, 7],
                       [15,  2,  2,  4,  8],
                       [ 6, 17,  6,  7,  8],
                       [ 1,  6, 12,  6,  2],
                       [ 8,  1,  7,  2, 10]]]

I have also played around with the number of iterations and it always seems like the gradient is slightly off in terms of the positions.