From 9427b269d9562964eb79593101fe29a7314c293e Mon Sep 17 00:00:00 2001 From: Junsong Chen Date: Mon, 20 Jan 2025 15:52:28 +0800 Subject: [PATCH] fix the bug for `AutoencoderDC` encode training; (#158) Signed-off-by: lawrence-cj --- diffusion/model/builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffusion/model/builder.py b/diffusion/model/builder.py index 89c7822..82e1fd9 100755 --- 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")