Hello all,
I  have a 3D tensor called x of size (batch_size, sequence_len, features).
For each batch element I want to take 3 rows, whose indexes could be different from batch element to batch element.
With a classic for loop over batch dimension it would look like:
res = torch.empty(x.size(0),3,x.size(2)) # take 3 rows for each batch
for i in range(x.size(0)):
    res[i] = x[i][ rows[i] ] # select appropriate rows

where rows is a LongTensor containing the indexes of the 3 rows to be taken for each batch element.
Is there a way to get rid of the for loop? I tried many ways but I am not able to make it work for a 3D tensor.
Thank you!