Torch filter multidimensional tensor by start and end values

I have a list of sentences and I am looking to extract contents between two items.
If the start or end item does not exist, I want it to return a row with padding only.
I already have the sentences tokenized and padded with 0 to a fixed length.

I figured a way to do this using for loops, but it is extremely slow, so would like to
know what is the best way to solve this.

import torch
start_value, end_value = 4,9

data = torch.tensor([
[3,4,7,8,10,0,0,0,0,0], # does not contain end value
[3,7,5,9,2,0,0,0,0,0], # does not contain start value

#expected output

There are two possible solutions

  1. Use the “for loop” then store the processed data in a file. This is to avoid reprocessing the data.

  2. You can use Pytorch dataloader to process your inputs in batches. This can save you so much time. Check Pytorch Dataloaders


i am doing this calculation inside the compute loss function, where the data is the prediction of the model, so I dont think I can use dataloaders here. In my for loop solution I used pad_sequence.
I was wondering if there is a way to do tensor operations instead of using for loop.