You could use index_select with repeated indices.
It’s definitely not the most elegant solution, but maybe it’s a bit easier to use.
x = torch.Tensor([[0, 1, 2], [3, 4 ,5]])
indices = torch.LongTensor([0, 0, 0, 1, 1, 1, 2, 2, 2])
torch.index_select(x, 1, indices)

Let me know, if this was helpful or still a painful solution