Hi. I’m trying to slice a tensor of different sizes for each batch. I have two tensors which are including start and end index respectively. For example, for a tensor of size [3, 1, 8], I want only specific tensors having different range:
start = torch.tensor([1, 3, 1])
end = torch.tensor([7, 7, 5])

# Do slicing

tensor([[[ 0,  1,  2,  3,  4,  5,  6,  7]], # [1, 7]

        [[ 8,  9, 10, 11, 12, 13, 14, 15]], # [3, 7]

        [[16, 17, 18, 19, 20, 21, 22, 23]]]) # [1, 5]

So, the first one will be [1,  2,  3,  4,  5,  6], the second one will be [11, 12, 13, 14] and the last one will be [17, 18, 19, 20]. I used iteration, but it is too slow and not elegant. How can I get them more faster and more elegant than just iterative way? The below shows my implementation.
context = torch.zeros(batch_size, 1, self.hid_dim).to(device)

for i in range(batch_size):
  local = torch.tensor([j for j in range(start[i], end[i])]).to(device) # [2D + 1]: for j
  batch_Pt = Pt[i]
  local = torch.exp(-1 * ((local-batch_Pt) ** 2)/((self.D ** 2) / 2)) # [2D + 1]
  score = align_score[i] * local  # [1, 2D + 1]
  context[i] = score.mm(H[i])