For a binary classification use case, you can use either one or two output channels.
For a single channel output, you could use nn.BCEWithLogitsLoss
.
If you are using two output channels, you could treat your use case as a
- multi-class classification (only one valid class per pixel) and use
nn.CrossEntropyLoss
- or as a multi-label classfication (zero, one or more classes valid per pixel), in which case you would again use
nn.BCEWithLogitsLoss
The background would also represent a class as answered in your other topic.