only square images for now

lucidrains · Jun 18, 2024 · 081d4cf · 081d4cf
1 parent 82ea31e
commit 081d4cf
Show file tree

Hide file tree

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

diff --git a/titok_pytorch/titok.py b/titok_pytorch/titok.py
@@ -61,6 +61,8 @@ def __init__(
         l - latents
         """
 
+        self.image_size = image_size
+
         assert divisible_by(image_size, patch_size)
 
         dim_patch = channels * patch_size ** 2
@@ -139,6 +141,8 @@ def forward(
         return_codebook_ids = False,
         return_recons = False
     ):
+        assert images.ndim == 4 and images.shape[-2:] == ((self.image_size,) * 2)
+
         batch = images.shape[0]
         orig_images = images