Solved by ptrblck in post #2 
              
              
                This should work: 
result2 = torch.zeros(3, 4)
result2[:, 1:] = x[y[:, :-1], y[:, 1:]]
print((result==result2).all())
> tensor(True)