[Python] Selecting at different column index for each row in tensor

I have a pytorch tensor

t = torch.tensor(
[[1.0, 1.5, 0.5, 2.0],
[5.0, 3.0, 4.5, 5.5],
[0.5, 1.0, 3.0, 2.0]]
)


t[:, [-1]] gives me last column value of each row:

tensor([[2.0000],
[5.5000],
[2.0000]])


However, I want to slice values at different columns per row. For example, in t for the 1st, 2nd and 3rd row, I want to slice at 2, -1, 0 index respectively to get the following tensor:

tensor([[0.5],
[5.5],
[0.5]])


How can I do it in torch?

Continue reading...