I’d probably take @ptrblck’s sampler and

compute the class weights in __init__,
have a tensor that contains the class weight for each index,
use torch.multinomial with the weights instead of randperm.

For the latter, you can compare RandomSampler 1
    def __iter__(self):
        return (self.indices[i] for i in torch.randperm(len(self.indices), generator=self.generator))

to WeightedRandomSampler 4
    def __iter__(self):
        rand_tensor = torch.multinomial(self.weights, self.num_samples, self.replacement, generator=self.generator)
        return iter(rand_tensor.tolist())

(except you want the indices her, too).
Note that you likely want replacement=True.
Best regards
Thomas