Dealing with imbalanced datasets in pytorch

[Skinish]

I am using Adam so it should not make much difference