I haven’t really made any changes to the code. I tried testing it out with BCEWithLogitsLoss() (which had also giveen pretty good results previously) in case there are any issues with just the Dice Loss but the loss isn’t really decreasing as such and it’s overall stuck at pretty low iou scores.