Update dependencyvit.py

huggingface · Dec 8, 2024 · a5b0f0d · a5b0f0d
1 parent 09e61fd
commit a5b0f0d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/timm/models/dependencyvit.py b/timm/models/dependencyvit.py
@@ -272,7 +272,7 @@ def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self._pos_embed(x)
         x = self.patch_drop(x)
         x = self.norm_pre(x)
-        x = x.reshape(B, -1, *self.patch_embed.dynamic_feat_size((H, W))) # [B, N, C] -> [B, C, H, W]
+        x = x.transpose(1, 2).reshape(B, -1, *self.patch_embed.dynamic_feat_size((H, W))) # [B, N, C] -> [B, C, H, W]
         m = torch.Tensor([1]).to(x)
         if self.grad_checkpointing and not torch.jit.is_scripting():
             x, m = checkpoint_seq(self.blocks, (x, m))