fix bug with einops

lucidrains · Mar 4, 2021 · 3088a2d · 3088a2d
1 parent c421f9f
commit 3088a2d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 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.4',
   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,7 +97,7 @@ 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),
+            Rearrange('b c (h p1) (w p2) -> (b h w) c p1 p2', p1 = patch_size, p2 = patch_size),
             nn.Unfold(pixel_width, stride = pixel_width),
             Rearrange('... c n -> ... n c'),
             nn.Linear(3 * pixel_width ** 2, pixel_dim)