diff --git a/timm/models/davit.py b/timm/models/davit.py index 97f8813b..713bc87e 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -164,11 +164,13 @@ class PatchEmbed(nn.Module): B, C, H, W = x.shape if self.norm.normalized_shape[0] == self.in_chans: x = self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - - x = F.pad(x, (0, self.patch_size[1] - torch.floor(W % self.patch_size[1]))) + print(self.patch_size[1] - W % self.patch_size[1]) + print(self.patch_size[0] - H % self.patch_size[0]) + + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) - x = F.pad(x, (0, 0, 0, self.patch_size[0] - torch.floor(H % self.patch_size[0]))) + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) x = self.proj(x)