I have a 2d Tensor A of shape (d, d) and I want to get the indices of its maximal element.
torch.argmax only returns a single index. For now I got the result doing the following, but it seems involuted.
vals, row_idx = A.max(0)
col_idx = vals.argmax(0)

And then, A[row_idx, col_idx] is the correct maximal value. Is there any more straightforward way to get this?