address #1

lucidrains · Jun 19, 2024 · 35ae3b5 · 35ae3b5
1 parent 8e56054
commit 35ae3b5
Show file tree

Hide file tree

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

diff --git a/titok_pytorch/titok.py b/titok_pytorch/titok.py
@@ -182,7 +182,7 @@ def forward(
         if return_codebook_ids:
             return indices
 
-        recon_images = self.decode(latents)
+        recon_images = self.decode(quantized)
 
         # reconstruction loss