Fix positional embedding resampling for non-square inputs in ViT

Author: wojtke

timm/models/vision_transformer.py

@@ -669,9 +669,11 @@ def _pos_embed(self, x: torch.Tensor) -> torch.Tensor:
 
         if self.dynamic_img_size:
             B, H, W, C = x.shape
+            prev_grid_size = self.patch_embed.grid_size
             pos_embed = resample_abs_pos_embed(
                 self.pos_embed,
-                (H, W),
+                new_size=(H, W),
+                old_size=prev_grid_size,
                 num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens,
             )
             x = x.view(B, -1, C)