Thanks for the advice. However, I’m confused with .diagonal() usage here
for example:
>>> x
 tensor([[[1,  2],
          [ 3, 4]],

          [[5, 6],
          [7,  8]]])

  >>> torch.diagona(x)

tensor([[1., 7.],
        [2., 8.]])


while I’m expecting with  [[1,4],[5,8]]. Am I missing anything?