torch.narrow() with tensor of start indices

I'm wondering if it's possible to do something in PyTorch...

Basically, what I'm trying to do is something like torch.narrow(), but where the start value is a tensor of indices.

Suppose I had the following tensors:

start_indices = [3, 1, 2, ...] # shape: (M,)
dataset = [
    3, 5, 3, 4, 8, 0, 1,
    3, 9, 7, 2, 7, 3, 7,
    6, 0, 2, 3, 0, 2, 5,
    ...
] # shape: (M, N) (here, N = 7, just for demo)

In my use case, M ~ 10^3 and N > 10^6

I'm wondering if it's possible to do something like this:

torch.narrow(
    input=dataset,
    dim=1,
    start=start_indices,
    length=4
)

that returns a "narrowed" version of the original dataset, where the start of each row is selected by its corresponding start index. To complete the example I gave above, the result would look like:

result = [
    4, 8, 0, 1,
    9, 7, 2, 7,
    2, 3, 0, 2,
    ....
] # shape: (M, 4)

I realize that depending on the start indices and length, you could have some lines going out of bounds of the original tensor -> I'm fine with receiving an error in this case.

Right now, I'm currently using a for-loop to do this, but it becomes far too slow as M becomes large. I'm hoping there's a way to do this that takes advantage of data parallelism.

Any ideas? I thought of using torch.index_select(), but it only selects within rows based off of single indices. I want to select within rows by single indices, but then extend those selections out by some fixed length. In that way, what I'm looking for is kind of a "combination" of torch.narrow() and torch.index_select().

Thanks in advance!