import torch


def calc_total_loss(flow, data):
    """
    Compute total loss of model on data.

    Parameters
    ----------
    flow : nflows.flows.Flow
        Normalising flow, instance of Flow object from nflows.flows, generated
        from e.g. load_flow() or setup_MAF().
    data : torch.Tensor, shape (N, n_dim)
        Torch tensor containing data on which to evaluate flow loss. Tensor
        should have shape (N, n_dim), where N is number of data points and
        n_dim is number of input dims of flow.

    Returns
    -------
    loss : float
        Model loss evaluated on given dataset.

    """
    flow.eval()
    with torch.no_grad():
        loss = -flow.log_prob(inputs=data, context=None).mean().item()
    return loss