fix the bug for AutoencoderDC encode training; (#158)

Signed-off-by: lawrence-cj <[email protected]>
NVlabs · Jan 20, 2025 · 9427b26 · 9427b26
1 parent 61b25ea
commit 9427b26
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/diffusion/model/builder.py b/diffusion/model/builder.py
@@ -113,7 +113,7 @@ def vae_encode(name, vae, images, sample_posterior, device):
     elif "AutoencoderDC" in name:
         ae = vae
         scaling_factor = ae.config.scaling_factor if ae.config.scaling_factor else 0.41407
-        z = ae.encode(images.to(device))
+        z = ae.encode(images.to(device))[0]
         z = z * scaling_factor
     else:
         print("error load vae")