How to tackle the class imbalance problem during training in PyTorch

2

You could over/undersample the classes using a WeightedRandomSampler.
Have a look at this post for a small example.