Ever considered implementing class balanced loss? I found your repo useful and would like to see it include more. https://github.com/vandit15/Class-balanced-loss-pytorch