select sub elements from another batch
I have first tensor with size  torch.Size([12, 64, 8, 8, 3]) ,  (8,8,3) is the image size, 64 is the patches, 12 is the batch size
There is another tensor with size torch.Size([12, 10]) which select 10 patches for each item in the batch (select 10 patches from total of 64). Thus it store the indexes. How to use this to query for the first tensor  with list comprehension?

You can use index_select:
c = [torch.index_select(i, dim=0, index=j) for i, j in zip(a,b)]

a and b are your tensor and indices respectively.
You could stack it in the zero dimension afterwards.
