From 825a9484d165065a544032b729ef00e5654afe8c Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Sun, 4 Oct 2020 12:39:51 -0700 Subject: [PATCH] small bug fix --- setup.py | 2 +- vit_pytorch/vit_pytorch.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index b976b18..b8d4755 100644 --- 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 index d63d648..47fd990 100644 --- 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