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.
Why don’t you:
loss = nn.CrossEntropyLoss2d(output[mask>0], target[mask>0])