How to efficiently retrieve the indices of maximum values in a Torch tensor? : <p>Assume to have a torch tensor, for example of the following shape:</p>

<pre><code>x = torch.rand(20, 1, 120, 120)
</code></pre>

<p>What I would like now, is to get the indices of the maximum values of each 120x120 matrix. To simplify the problem I would first <code>x.squeeze()</code> to work with shape <code>[20, 120, 120]</code>. I would then like to get torch tensor which is a list of indices with shape <code>[20, 2]</code>.</p>

<p>How can I do this fast?</p>
