Solved by ptrblck in post #2 
              
              
                I assume your labels are saved in a one-hot encoded format for a multi label classification. 
If that’s the case, you could iterate your Dataset once and just count all class occurrences: 
class MyDataset(Dataset):
    def __init__(self, num_classes):
        self.data = torch.randn(100, 2)
        …