title: Top K indices of a multi-dimensional tensor
body: <p>I have a 2D tensor and I want to get the indices of the top k values. I know about <a href="https://pytorch.org/docs/stable/generated/torch.topk.html" rel="nofollow noreferrer">pytorch's topk</a> function. The problem with pytorch's topk function is, it computes the topk values over some dimension. I want to get topk values over both dimensions.</p>
  <p>For example for the following tensor</p>
  <pre class="lang-py prettyprint-override"><code>a = torch.tensor([[4, 9, 7, 4, 0],
          [8, 1, 3, 1, 0],
          [9, 8, 4, 4, 8],
          [0, 9, 4, 7, 8],
          [8, 8, 0, 1, 4]])
  
  </code></pre>
  <p>pytorch's topk function will give me the following.</p>
  <pre class="lang-py prettyprint-override"><code>values, indices = torch.topk(a, 3)
  
  print(indices)
  # tensor([[1, 2, 0],
  #        [0, 2, 1],
  #        [0, 1, 4],
  #        [1, 4, 3],
  #        [1, 0, 4]])
  
  </code></pre>
  <p>But I want to get the following</p>
  <pre><code>tensor([[0, 1],
          [2, 0],
          [3, 1]])
  
  </code></pre>
  <p>This is the indices of 9 in the 2D tensor.</p>
  <p>Is there any approach to achieve this using pytorch?</p>