def dataset_mean_loss(trainer, data_loader, device):
    """Returns mean loss of model on a dataset. Useful for calculating
    validation loss.

    Parameters
    ----------
    trainer : training.Trainer instance
        Trainer instance for model we want to evaluate.

    data_loader : torch.utils.data.DataLoader

    device : torch.device
    """
    epoch_loss = 0.
    for x_batch, y_batch in data_loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        y_pred = trainer.model(x_batch)
        loss = trainer._loss(y_pred, y_batch)
        epoch_loss += loss.item()
    return epoch_loss / len(data_loader)