I think this fixes the bug where vq loss is small compared to recon l…

…oss (#125)
clementchadebec · Jan 21, 2024 · b5c7462 · b5c7462
1 parent 18fe279
commit b5c7462
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/pythae/models/vq_vae/vq_vae_utils.py b/src/pythae/models/vq_vae/vq_vae_utils.py
@@ -50,14 +50,14 @@ def forward(self, z: torch.Tensor, uses_ddp: bool = False):
         commitment_loss = F.mse_loss(
             quantized.detach().reshape(-1, self.embedding_dim),
             z.reshape(-1, self.embedding_dim),
-            reduction="mean",
+            reduction="sum",
         )
 
         embedding_loss = F.mse_loss(
             quantized.reshape(-1, self.embedding_dim),
             z.detach().reshape(-1, self.embedding_dim),
-            reduction="mean",
-        ).mean(dim=-1)
+            reduction="sum",
+        )
 
         quantized = z + (quantized - z).detach()
 
@@ -147,7 +147,7 @@ def forward(self, z: torch.Tensor, uses_ddp: bool = False):
         commitment_loss = F.mse_loss(
             quantized.detach().reshape(-1, self.embedding_dim),
             z.reshape(-1, self.embedding_dim),
-            reduction="mean",
+            reduction="sum",
         )
 
         quantized = z + (quantized - z).detach()