small bug fix

lucidrains · Oct 4, 2020 · 825a948 · 825a948
1 parent ee8088b
commit 825a948
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 = 'vit-pytorch',
   packages = find_packages(),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   author = 'Phil Wang',

diff --git a/vit_pytorch/vit_pytorch.py b/vit_pytorch/vit_pytorch.py
@@ -86,7 +86,7 @@ def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, ml
     def forward(self, img):
         p = self.patch_size
 
-        x = rearrange(img, 'b c (p1 h) (p2 w) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
+        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
         x = self.patch_to_embedding(x)
         x = torch.cat((self.cls_token, x), dim=1)
         x += self.pos_embedding