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 ?