readme

lucidrains · Jun 19, 2024 · 6dcc12b · 6dcc12b
1 parent 166a38f
commit 6dcc12b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -30,7 +30,7 @@ loss.backward()
 # after much training
 # extract codes for gpt, maskgit, whatever
 
-codes = titok.tokenize(images)
+codes = titok.tokenize(images) # (2, 32)
 
 # reconstructing images from codes
 

diff --git a/tests/test_titok.py b/tests/test_titok.py
@@ -6,7 +6,10 @@ def test_titok():
 
     images = torch.randn(2, 3, 256, 256)
 
-    titok = TiTokTokenizer(dim = 512)
+    titok = TiTokTokenizer(
+        dim = 512,
+        num_latent_tokens = 32
+    )
 
     loss = titok(images)
     loss.backward()
@@ -16,6 +19,8 @@ def test_titok():
 
     codes = titok.tokenize(images)
 
+    assert codes.shape == (2, 32)
+
     # reconstructing images from codes
 
     recon_images = titok.codebook_ids_to_images(codes)