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])