Working with argmax
In higher dimensions torch.argmax method returns the list of indices of maximum values according to the specified axis. we can understand it with an example of how argmax() method works on 2 Dimensional tensors.
Example: [[1,10], [20,15]]
Two Dimensional tensor has only two axis-0,1 (Rows & Columns)
- Along axis-0 The argmax method check for the maximum value in each column and returns the one index position of maximum value in each column. So in first column 20 is the maximum value and index value of it is 1 and in second column 15 is the maximum value and it’s Index is 1. so argmax method returns [1,1] as result.
- Along the axis-1 The argmax method check for the maximum value in each row and returns the one index position of maximum value in each row. So in first row 10 is the maximum value and index value of it is 1 and in second row 20 is the maximum value it’s index position is 0. so argmax method returns [1,0] as result.
How Does torch.argmax Work for 4-Dimensions in Pytorch
In this article, we are going to discuss how does the torch.argmax work for 4-Dimensions with detailed examples.