Solved by ptrblck in post #2 
              
              
                torch.gather should work fine: 
b = torch.tensor([4, 5, 6, 7, 8], device='cuda')
c = torch.tensor([0, 0, 0, 1, 1, 3, 4], device='cuda')

res = torch.gather(b, 0, c)
print(res)
> tensor([4, 4, 4, 5, 5, 7, 8], device='cuda:0')