I see, just integrated the proposed method on another topic here to use WeightedRandomSampler and will post the result here when the training and test phases are completed. Between, could you recommend any utility function to check the class distribution and the prediction for unique values?