copperwiring:

The end goal is segmentation in an image i.e. I have two classes only and each pixel/feature in the image can belong to one class.


In that case you could either:

Use nn.BCEWithLogitsLoss as your criterion, make sure your model outputs logits with the shape [batch_size, 1, height, width] and use a target with the same shape containing your labels (0 and 1). For the accuracy calculation, you could apply a threshold of 0 to get the predicted class as: preds = outputs > 0.

Treat your binary segmentation as a multi-class segementation. Use nn.CrossEntropyLoss as the criterion, make sure your model outputs logits in the shape [batch_size, nb_classes, height, width] and your target is a LongTensor in the shape [batch_size, height, width] containing the class indices (0 and 1). To get the prediction, you could stick to the torch.argmax approach.