How to slice 2D Torch tensor individually per row?
I have a 2D tensor in Pytorch that I would like to slice:
x = torch.rand((3, 5))

In this example, the tensor has 3 rows and I want to slice x, creating a new tensor y that also has 3 rows and num_col cols.
What's challenging for me is that I want to slice different columns per row. All I have is x, num_cols, and idx, which is a tensor holding the start index from where to slice.
Example:
What I have is num_cols=2, idx=[1,2,3] and
x=torch.arange(15).reshape((3,-1)) =
tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]])

What I want is
y=
tensor([[ 1,  2],
        [ 7,  8],
        [13, 14]])

What's the &quot;torch&quot;-way of doing this? I know, I can slice if I get a boolean mask somehow, but I don't know how to construct that with idx and num_cols without normal Python loops.

You could use fancy indexing together with broadcasting. Another solution might be to use torch.gather which is similar to numpy's take_along_axis. Your idx array would need to be extended with the extra column:
x = torch.arange(15).reshape(3,-1)
idx = torch.tensor([1,2,3])

idx = torch.column_stack([idx, idx+1])
torch.gather(x, 1, idx)

output:
tensor([[ 1,  2],
        [ 7,  8],
        [13, 14]])

