Hello,

I am trying to do a binary segmentation task on the coco dataset using RestNet18 as encoder and DeepLabV3 as the decoder. Looking around i found that BCEWithLogitsLoss is the recommended criterion but i cannot get my training to work. Can you please help me proceed in the right direction ?

I am trying the following in the train method :

outputs = model(inputs)

loss = criterion(torch.nn.Sigmoid(outputs), labels)

The input to my network is torch.Size([32, 3, 256, 256])

The Output is torch.Size([32, 2, 256, 256])

and the labels are torch.Size([32, 1, 256, 256])

I cannot get the sigmoid to work, earlier i tried the CrossEntropyLoss with 2 classes but i doubt if it even worked. I am mostly stressed about the shapes.

Will be grateful for any help, thanks !