My segmentation model having an issue with imbalanced data. I want to use class weight parameters just like in another framework. What is the alternative case in Pytorch? I’m using dice loss should I change the loss?