Exploding KL loss

Hello. Thank you for providing this.
I am trying to train the model and the KL loss seems to explode for some of the steps. I am wondering if you faced a similar issue? 

<img width="401" alt="image" src="https://github.com/user-attachments/assets/f5320332-8cfc-4aa8-9150-c081560cc74a">
