# Why did my customized KL divergence cause NAN?

**URL:** <https://discuss.pytorch.org/t/why-did-my-customized-kl-divergence-cause-nan/214339>\
**Category:** vision\
**Created:** [December 18, 2024, 12:23am UTC](https://discuss.pytorch.org/t/why-did-my-customized-kl-divergence-cause-nan/214339 "2024-12-18T00:23:20Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![JimW](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jimw/32/38000_2.png) [@JimW](https://discuss.pytorch.org/u/JimW)\
**Post date:** [December 18, 2024, 12:23am UTC](https://discuss.pytorch.org/t/why-did-my-customized-kl-divergence-cause-nan/214339/1 "2024-12-18T00:23:20Z")

</div>

According to some constraint within our project, I need to rewrite KL divergence computation with basic PyTorch operations. Here is my implementation of KL divergence.  
`torch.mean(q*torch.log(q/p)`  
I compared the output values of this implement with the output values from the following:  
`nn.KLDivLoss()(p.log(), q)`  
And they generate the same values for my limited test set.

When I kept using

> nn.KLDivLoss()(p.log(), q)

, everything was fine. However, when I replaced it with

> torch.mean(q\*torch.log(q/p)

Some part of my model started generating NAN after several training epochs.

Could someone provide some guidance how this could happen and how to fix it?

---

<div class="post-metadata">

**Author:** ![KFrank](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/32.png) [@KFrank](https://discuss.pytorch.org/u/KFrank)\
**Post date:** [December 22, 2024, 3:13am UTC](https://discuss.pytorch.org/t/why-did-my-customized-kl-divergence-cause-nan/214339/2 "2024-12-22T03:13:56Z")

</div>

Hi Jim!

> [@JimW](#):
>
> > torch.mean(q\*torch.log(q/p)
> 
> Some part of my model started generating NAN after several training epochs.

Depending on your use case, computing `log (q) - log (p)` can be more stable  
numerically than computing `log (q / p)`. (They are the same mathematically,  
but can differ numerically.)

Try using `torch.mean (q * (q.log() - p.loq()))`as your KL divergence  
implementation.

As an aside, `p` would typically be a prediction, that is, the output of your model.  
If so – for example, if the last layer of your model is a `Linear` – the output of  
your model – let’s call it `y` – might naturally already be in log-space. That is,  
your model might naturally be outputting `y = p.log()`. If so, just use `y` directly  
(without ever explicitly converting it to `p = y.exp()`). In this case, your KL  
divergence would become `torch.mean (q * (q.log() - y))`.

Best.

K. Frank