A custom sampler sounds like the right idea.
Once you have the labels, you could create a custom sampler (e.g. as a BatchSampler) and let it pass the indices for a single target only. Here is an example of the usage of a BatchSampler and for your use case you could derive from this class to implement a custom sample logic.
1 Like