title: How to efficiently retrieve the indices of maximum values in a Torch tensor?
body: <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>