Output problem with torch.nn.CrossEntropyLoss()

Hello. I’m doing semantic segmentation with Unet. My target class of 0 (background) and 1 (class-of-interest) is in one channel. So, my target is [batch_size, 256, 256]. And, my model output has two channels (one for each class) [batch_size, 2, 256, 256]. I want to use torch.nn.CrossEntropyLoss() for my loss function but I don’t know how to properly reduce the output of my model to one channel after forward pass so that I could insert it in the loss function. Any help is appreciated.

I think [B, C, 256, 256] is correct to apply nn.CrossEntropyLoss.
see how it works CrossEntropyLoss — PyTorch 1.10 documentation

1 Like

Thank you. You’re right. I had a problem in my code and I thought the channel difference was the issue. The issue was something else. thanks again.