I’m working on an Inpainting project.
I’m trying to calculate the L2 reconstruction loss of only the regions that have been masked.
I can mask the output with
criterion = nn.MSELoss()
masked_pred = torch.mul(prediction, mask) 
loss = criterion(masked_input, masked_pred) 

But the problem is that MSELoss averages the loss over all pixels, and since most of my pixels are masked the loss becomes very small.
Is there a way to calculate the MSELoss only on non-masked pixels?
Just as a side note, the mask that I use is a randomly generated mask (not fixed)
Thanks in advance