How to train with imbalance class for classification?

I think your issue might be related to this one.
Try to pass the weights for each sample to WeightedRandomSampler and see if it’s working.