able to return reconstructed images with loss

lucidrains · Jun 18, 2024 · 8e56054 · 8e56054
1 parent d14e6f1
commit 8e56054
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "titok-pytorch"
-version = "0.0.2"
+version = "0.0.3"
 description = "TiTok - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

diff --git a/titok_pytorch/titok.py b/titok_pytorch/titok.py
@@ -139,7 +139,7 @@ def forward(
         self,
         images,
         return_codebook_ids = False,
-        return_recons = False
+        return_recon_images = False
     ):
         assert images.ndim == 4 and images.shape[-2:] == ((self.image_size,) * 2)
 
@@ -182,10 +182,13 @@ def forward(
         if return_codebook_ids:
             return indices
 
-        recon = self.decode(latents)
+        recon_images = self.decode(latents)
 
         # reconstruction loss
 
-        recon_loss = F.mse_loss(recon, orig_images)
+        recon_loss = F.mse_loss(recon_images, orig_images)
 
-        return recon_loss
+        if not return_recon_images:
+            return recon_loss
+
+        return recon_loss, recon_images