WeightedRandomSampler for dataset balance

weights should be given as sample weights (i.e. each sample in your Dataset will get a weight) not as class weights. Have a look at this small code snippet to see an example.