Hi,
I wanted to try your code on SWIMSEG dataset. It has only two classes- cloud and sky (background). I got IoU of 80%. But I wanted to BCELoss instead of CrossEntropyLoss so that I could reduce the number of output channels to just 1 instead of 2. Hence, I used use BCELoss (For this, I also introduced a sigmoid layer at the end of the model). The problem is that it is now giving huge error and output tensor is quickly moving towards 0, i.e., I am getting all the outputs as 0s. I tried changing learning rates. But it was not very helpful. Any idea, what's going wrong?