Hey, sorry for a late reply.
I didn’t really understand your solution, sampler is not a function, but an iterable (you can think of it as a stream of indices). I think that one way to impement what you want would be to:

Use vanilla DataLoader with batch_size=num_class.
When your dataset is asked for ith sample, you look up its class and sample num_exp_per_class-1 other examples from that class, concatenate them into a “sub-batch” and return that.
The data loader will return num_class x num_exp_per_class x <feature dims> elements.