Hi, Let’s say we have an input tensor of shape C x H x W
, and an index tensor of shape N x 2 x 2
, where N
is the number of sub tensors I want to extract from input tensor, each index[i, :, :]
is a 2 x 2 matrix where the 1st row specifies the starting point and distance along dim H
, and the 2nd row along dim W
, a minimal working example is as follows:
input = torch.Tensor(C, H, W)
index = torch.LongTensor(N, 2, 2)
subs = []
for i in range(N):
h_s, h_dis = index[i, 0, 0], index[i, 0, 1]
w_s, w_dis = index[i, 1, 0], index[i, 1, 1]
subs.append(input[:, h_s: h_s + h_dis, w_s: w_s + w_dis])
I’d like this indexing to be paralleled so that the for loop
can be eliminated because I’ll do some compute-intensive stuff inside the for loop
after I extract a sub tensor, any idea on how to do it? Besides, each sub tensor can be of varying size, so I need a list to store sub tensors, instead of concatenation since it requires dimension should be equal.