Is there an alternative to using torch.diagonal() which does not alter the ordering of dimensions of a tensor?
For example, I have a 4D tensor, A, of shape (2,4,4,6). I want B = A[:,i,i,:] for all i in range(4).
I did the following in PyTorch:
tmp = torch.diagonal(A,dim1=1,dim2=2)
B = tmp.permute(0,2,1).contiguous()

Is there a way to avoid the permute() operation?
Many thanks.