I’m not sure I understand the use case correctly, but it sounds as if you are looking for the torch.utils.data.Subset class, which would allow you to pass indices to it and only return samples from these indices.
1 Like