theairbend3r:

So, how do we ensure that our train and val sets as well as batches in train_loader and val_loader have an equal distribution of output classes?


I would use a stratified split using e.g. sklearn.model_selection.train_test_split, calculate the weights separately for each split, and use a WeightedRandomSampler for each subset to get balanced batches.