Solved by eqy in post #5 
              
              
                Does torch.gather work for you? 
import torch
y = torch.arange(50).reshape(5,10)
s = torch.arange(10).reshape(5,2)
print(torch.gather(y, 1, s))

tensor([[ 0,  1],
        [12, 13],
        [24, 25],
        [36, 37],
        [48, 49]])