From 08d8ef6e7283d5d3a0aa17ead824e5ea255fff68 Mon Sep 17 00:00:00 2001
From: qqingzheng <2533221180@qq.com>
Date: Wed, 13 Mar 2024 18:16:28 +0800
Subject: [PATCH] [bug] fix quantization loss

Former-commit-id: 111fcf55f969ffe4c9be284a317ae5712570a527
---
 opensora/models/ae/videobase/vqvae/trainer_vqvae.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/opensora/models/ae/videobase/vqvae/trainer_vqvae.py b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py
index 5b465c7f6..419c18244 100644
--- a/opensora/models/ae/videobase/vqvae/trainer_vqvae.py
+++ b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py
@@ -15,8 +15,9 @@ class VQVAETrainer(Trainer):
         vq_output = model.codebook(z)
         x_recon = model.decoder(model.post_vq_conv(vq_output["embeddings"]))
         recon_loss = F.mse_loss(x_recon, x) / 0.06
-
-        return recon_loss
+        commitment_loss = vq_output['commitment_loss']
+        loss = recon_loss + commitment_loss
+        return loss
 
     def _save(self, output_dir: Optional[str] = None, state_dict=None):
         output_dir = output_dir if output_dir is not None else self.args.output_dir