import torch


def deen_loss(energy_model, x, sigma=0.1):
    """DEEN loss from
    Deep Energy Estimator Networks

    The loss is computed as
    x_ = x + v   # noisy samples
    s = -dE(x_)/dx_
    loss = 1/2*||x - x_ + sigma^2*s||^2

    Args:
        x (torch.Tensor): input samples
        sigma (int, optional): noise scale. Defaults to 1.

    Returns:
        DEEN loss
    """
    x = x.requires_grad_()
    v = torch.randn_like(x) * sigma
    x_ = x + v
    s = sigma ** 2 * energy_model.score(x_)
    loss = torch.norm(s + v, dim=-1) ** 2
    loss = loss.mean() / 2.0
    return loss