import torch


def mse_masked(input, target, mask):
    """Mean squared error with visibility mask implementation
    Note:
        If mask is all zero => loss should be zero
    """
    return torch.mean(torch.sum((input - target) ** 2, dim=2) * (mask == 1))