[bug] fix quantization loss

Former-commit-id: 111fcf55f969ffe4c9be284a317ae5712570a527
2024-03-13 18:16:28 +08:00
parent 51581271b4
commit 08d8ef6e72
1 changed files with 3 additions and 2 deletions
@@ -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