From 420d3b7588a6120d7fe6ce14b917e62664958d3b Mon Sep 17 00:00:00 2001
From: lawrence-cj <cjs1020440147@icloud.com>
Date: Mon, 20 Jan 2025 15:49:14 +0800
Subject: [PATCH] fix the bug for `AutoencoderDC` encode training;
Signed-off-by: lawrence-cj <cjs1020440147@icloud.com>
---
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")