Mask for CrossEntropyLoss2d

is it possible to only calculate the loss within mask area?

You can process the output a bit or write your own loss function to do that.

I implemented a weighted NLLLoss2d in this thread. Maybe you could use it by setting your weights accordingly to 0 and 1. :slight_smile:

1 Like

Why don’t you:

loss = nn.CrossEntropyLoss2d(output[mask>0], target[mask>0])

?