Balanced batch sampling with DistributedSampler/DDP

This post explaining an approach to create a DistributedWeightedSampler might be helpful.