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()