1 files changed, 2 insertions, 2 deletions

diff --git a/training/functional.py b/training/functional.py
index 8f47734..ccbb4ad 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -261,8 +261,8 @@ def loss_step(
     eval: bool = False
 ):
     # Convert images to latent space
-    latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach()
-    latents = latents * 0.18215
+    latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
+    latents = latents * vae.config.scaling_factor
 
     generator = torch.Generator(device=latents.device).manual_seed(seed + step) if eval else None