From 1c70271778dbccb488e32b39ed79dc8fbccd2cfc Mon Sep 17 00:00:00 2001
From: Yonghye Kwon <developer.0hye@gmail.com>
Date: Sat, 21 Aug 2021 22:25:46 +0900
Subject: [PATCH] Support image with width and height less than the image_size

Related to #145
---
 vit_pytorch/pit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vit_pytorch/pit.py b/vit_pytorch/pit.py
index 166f8e1..a5f92f8 100644
--- a/vit_pytorch/pit.py
+++ b/vit_pytorch/pit.py
@@ -175,7 +175,7 @@ class PiT(nn.Module):
 
         cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
         x = torch.cat((cls_tokens, x), dim=1)
-        x += self.pos_embedding
+        x += self.pos_embedding[:, :n+1]
         x = self.dropout(x)
 
         x = self.layers(x)