oh if you have class imbalance, use a weighted sampler, so that you see all classes with equal probability. That should help a lot (depending on the severity of your imbalance):
http://pytorch.org/docs/data.html#torch.utils.data.sampler.WeightedRandomSampler
You can give this to your DataLoader