fix pixel calculation bugs

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

Hide file tree

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

diff --git a/transformer_in_transformer/tnt.py b/transformer_in_transformer/tnt.py
@@ -97,10 +97,10 @@ def __init__(
         self.patch_tokens = nn.Parameter(torch.randn(num_patch_tokens + 1, patch_dim))
 
         self.to_pixel_tokens = nn.Sequential(
-            Rearrange('b c (p1 h) (p2 w) -> (b h w) c p1 p2', p1 = patch_size, p2 = patch_size),
-            nn.Unfold(pixel_width, stride = pixel_width),
+            Rearrange('b c (h p1) (w p2) -> (b h w) c p1 p2', p1 = patch_size, p2 = patch_size),
+            nn.Unfold(pixel_size, stride = pixel_size),
             Rearrange('... c n -> ... n c'),
-            nn.Linear(3 * pixel_width ** 2, pixel_dim)
+            nn.Linear(3 * pixel_size ** 2, pixel_dim)
         )
 
         self.patch_pos_emb = nn.Parameter(torch.randn(num_patch_tokens + 1, patch_dim))