As a continuation of this question,

  
    
    
    How to find input corresponding to argmin of output? 
  
  
    I have a basic scenario like this, 
import torch

def someFn(w):
    
    return w.sum(dim=1)

w = torch.rand(10,3,2)
fW = someFn(w)

fMinW, fIndex = fW.min(dim=1)

# Select wmin of size 10,3

Now I want to select the 10 (1x3) vectors that correspond to the fMinW output. How can I achieve this ?
  


I have a scenario like this,
import torch

def someFn(w):
    
    return w.sum(dim=1)

w = torch.rand(10,3,6)
fW = someFn(w)

fIndex = fW.argsort(dim=-1)

# Select wmin of size 10,3,2

I want to select the minimum m=2 w that correspond to the two minimum output.
I tried
wmin = w[torch.arange(w.size(0)), :, fIndex[:, :m]]

But it throws error.