I wrote the loss class:
class MaskedMSELoss(torch.nn.Module):
    def __init__(self):
        super(MaskedMSELoss, self).__init__()

    def forward(self, input, target, mask):
        diff2 = (torch.flatten(input) - torch.flatten(target)) ** 2.0
        sum2 = 0.0
        num = 0

        flat_mask = torch.flatten(mask)
        assert(len(flat_mask) == len(diff2))
        for i in range(len(diff2)):
            if flat_mask[i] == 1:
                sum2 += diff2[i]
                num += 1

        return sum2 / num

What would be default backward looks like?