extra assert for safe measure

lucidrains · Mar 4, 2021 · 41543d8 · 41543d8
1 parent 7621ce6
commit 41543d8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'transformer-in-transformer',
   packages = find_packages(),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Transformer in Transformer - Pytorch',
   author = 'Phil Wang',

diff --git a/transformer_in_transformer/tnt.py b/transformer_in_transformer/tnt.py
@@ -87,6 +87,7 @@ def __init__(
     ):
         super().__init__()
         assert image_size % patch_size == 0, 'image size must be divisible by patch size'
+        assert patch_size % pixel_size == 0, 'patch size must be divisible by pixel size for now'
 
         num_patch_tokens = (image_size // patch_size) ** 2
         pixel_width = patch_size // pixel_size