Could you explain the use case of creating the bags a bit more, please?
If I understand it correctly, you would like to sample from specific indices to create these bags?
Would you like to return a data tensor in [batch_size, 10, ...] or would 10 already be the batch size?