RHVA example notebook crashes. Exploding norm during leapfrog

[Theaublanchard]

Bug description
During training of the RHVAE on the simplest case, the sampling ultimately yields latent sample of exploding magnitude. Thus the evaluation of the log density of the sample returns NaN values.

More precisely, the error occurs because the method _log_z from the RHVAE class is called with infinite arguments. Those come from the second step of the leapfrog integrator (_leap_step_2).
Upon logging the norm of the gradient of the Hamiltonian, it explodes randomly after some time. The same phenomenon appears when using the "bce" reconstruction loss, albeit the training is stable for longer.

To Reproduce

1. Run the example notebook at rhvae_training.ipynb
2. The training will fail after a few epochs (see error stack below). Decreasing the learning rate only delays the crash.

Expected behavior
No crash during training

Desktop

• OS: SUSE Linux Enterprise Server 15 SP6
• python: 3.12.4
• numpy: 1.26.4
• torch: 2.5.1

Possible fix
Clip the norm of the gradient on each call to the gradient of the Hamiltonian. It could disturb the training, this needs to be explored.