Fix positional embedding resampling for non-square inputs in ViT

timm/models/deit.py

@@ -75,9 +75,11 @@ def set_distilled_training(self, enable=True):
     def _pos_embed(self, x):
         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)

timm/models/eva.py

@@ -560,9 +560,11 @@ def _pos_embed(self, x) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         if self.dynamic_img_size:
             B, H, W, C = x.shape
             if self.pos_embed is not None:
+                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=self.num_prefix_tokens,
                 )
             else:

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)