import torch


def mse(y_pred, y_true, masks=None):
    """Compute mean square error (MSE) loss with masks.

    Parameters
    ----------
    y_pred : :obj:`torch.Tensor`
        predicted data
    y_true : :obj:`torch.Tensor`
        true data
    masks : :obj:`torch.Tensor`, optional
        binary mask that is the same size as `y_pred` and `y_true`; by placing 0 entries in the
        mask, the corresponding dimensions will not contribute to the loss term, and will therefore
        not contribute to parameter updates

    Returns
    -------
    :obj:`torch.Tensor`
        mean square error computed across all dimensions

    """
    if masks is not None:
        return torch.mean(((y_pred - y_true) ** 2) * masks)
    else:
        return torch.mean((y_pred - y_true) ** 2)