title: Rearranging a 3-D array using indices from sorting?
body: <p>I have a 3-D array of random numbers of size <code>[channels = 3, height = 10, width = 10]</code>.</p>
  
  <p><a href="https://i.stack.imgur.com/eMzx0.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/eMzx0.png" alt="enter image description here"></a>
  Then I sorted it using sort command from pytorch along the columns and obtained the indices as well.</p>
  
  <p><a href="https://i.stack.imgur.com/H3JYY.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/H3JYY.png" alt="enter image description here"></a></p>
  
  <p>The corresponding index is shown below:</p>
  
  <p><a href="https://i.stack.imgur.com/qdTkg.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/qdTkg.png" alt="enter image description here"></a></p>
  
  <p>Now, I would like to return to the original matrix using these indices. I currently use <code>for</code> loops to do this (without considering the batches). The code is:</p>
  
  <pre><code>import torch
  torch.manual_seed(1)
  ch = 3
  h = 10
  w = 10
  inp_unf = torch.randn(ch,h,w)
  inp_sort, indices  = torch.sort(inp_unf,1)
  resort = torch.zeros(inp_sort.shape)
  
  for i in range(ch):
      for j in range(inp_sort.shape[1]):
          for k in range (inp_sort.shape[2]):
              temp = inp_sort[i,j,k]
              resort[i,indices[i,j,k],k] = temp
  </code></pre>
  
  <p>I would like it to be vectorized considering batches as well i.e.input size is <code>[batch, channel, height, width]</code>.</p>