Awesome! that solves the problem, thanks G.M
Small followup tho. lets say you have
x = tc.randn(5,10,15,20)
indexes # shape of 5,10
so now the indexes tell me which guy from dim=2 is the best, but then this best guy has a dimenions of not 1, but 20. How can apply indexes to x now? In a for loop it would be like:
for i
    for j
        out = x[i,j,indexes[i,j],:]

The final shape of out is 5,10,20.