Trying to write equivalent Poisson loss function using torch distribution

1

I’m currently using PoissonNLLLoss (well actually F.poisson_nll_loss) but I wanted to check if I can write my own custom loss using the poisson distribution from torch.distributions:

def poisson_nll(obs, lambd):
    poisson_dist = dist.Poisson(lambd)
    poisson_prob = poisson_dist.log_prob(obs)
    nll = -poisson_prob.mean()
    return nll

Does this look corect?