Subsetting a two dimensional tensor with a one dimensional tensor
I want to extract from each row of a two-dimensional tensor the column that is stored in another one dimensional tensor.
import torch
test_tensor = tensor([1,-2,3], [-2,7,4]).float()
select_tensor = tensor([1,2])

So in this particular example I would like to get the element in position 1 for the first row (so -2) and the element in position 2 for the second row (so 4).
I tried:
test_tensor[:, select_tensor]

But this selects the elements at position 1 and 2 for each row. I suspect it might be something very simple that I am missing.

You can use torch.gather
import torch
test_tensor = torch.tensor([[1,-2,3], [-2,7,4]]).float()
select_tensor = torch.tensor([1,2], dtype=torch.int64).view(-1,1) # number of dimension should match with the test tensor.
final_tensor = torch.gather(test_tensor, 1, select_tensor)
final_tensor

output
tensor([[-2.],
        [ 4.]])

or, use torch.view to flatten the output tensor: final_tensor.view(-1) will give you tensor([-2.,  4.])
