import torch


def thesis_kl_div_add_mse_loss(p, q, lam=1):
    """
    calculate the sum of kl divergence and mse loss
    :param p: p in the formula (P20-2) output similarities
    :param q: q in the formula (P20-2) input similarities
    :param lam: the constant that balances the influence of two losses
    :return: torch.tensor of the shape (,)
    """
    return torch.sum(p * torch.log(p / q)) + lam * torch.sum((p - q) ** 2)