No Variance Scaling of Latent Space for Kl-Regulerized Autoencoder

[maltesilber]

In the paper it states that the latent space get scaled based on the variance of the first batch. See:

However this behaviour does not seem to be implemented, as this is the forward pass used in the training/validation step: https://github.com/CompVis/latent-diffusion/blob/a506df5756472e2ebaf9078affdde2c4f1502cd4/ldm/models/autoencoder.py#L335-L342

When trying to fine-tune my own model I released this as the only loss that worsens is the KL-Loss. In the end I got a model with extremely high variance in the latent space.

Is my understanding correct that the scaling logic is missing from the AutoEncoder?