Thank you for your answer. I realize that my example was too simple, because I need to do this on 3D tensors. My solution for now is:
x = torch.as_tensor(
        [
            [[1, 2, 3, 4, 5], [6, 7, 8, 9, 0], [11, 12, 13, 14, 15]],
            [[16, 17, 18, 19, 20], [21, 22, 23, 24, 25], [26, 27, 28, 29, 30]],
        ]
    )

index = [[0, 0, 1], [1, 2, 0]]
# tensor([[ 6,  7,  8,  9,  0],
#         [11, 12, 13, 14, 15],
#         [16, 17, 18, 19, 20]])
x[index]

index_t = torch.as_tensor(index)
# tensor([[ 6,  7,  8,  9,  0],
#         [11, 12, 13, 14, 15],
#         [16, 17, 18, 19, 20]])
x = x.index_select(0, index_t[0])
x = x[torch.arange(x.shape[0]).unsqueeze(-1), index_t[1].unsqueeze(-1)].squeeze()


It is a bit convoluted, but it gets the job done. If there is a better way to do it, I will be happy to learn