import torch


def _masked_loss_func(output, masks, samples):
    """ Computes the MSE loss.

    :param output: tensor(Float); imputed samples
    :param masks: tensor(Float); corresponding masks
    :param samples: tensor(Float); original samples (with missing values)
    :return: tensor(Float); loss obtained
    """
    if masks is None:
        masks = torch.zeros(samples.shape)
    mse_loss = torch.sum((~masks.bool() * samples - ~masks.bool() * output) ** 2)
    if torch.sum(1 - masks) > 0:
        mse_loss /= torch.sum(1 - masks)
    return mse_loss