Pytorch Indexing
I have a tensor [[1,2],[4,5],[7,8]] and a tensor with indices [0,1,0].

I want to apply them to second dimension so that it returns: [1,5,8].

How should I do that?

Thanks!

import torch

arr=torch.tensor([[1,2],[4,5],[7,8]])
indices_arr=torch.tensor([0,1,0])

ret=arr[[0,1,2],indices_arr]
# print(ret)
# tensor([1, 5, 7])

