WeightedRandomSampler only sampling from one class

[ptrblck]

WeightedRandomSampler expects a weight tensor assigning the weight values to each sample, not the class index.
Have a look at this post which gives you an example.