Update davit.py

3 years ago · 42f86c8c09
parent 87abcafaa5
commit 42f86c8c09
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -111,7 +111,7 @@ class PatchEmbed(nn.Module):
            self.norm = nn.LayerNorm(in_chans)
-    def forward(self, x, size: Tuple[int, int]):
+    def forward(self, x : Tensor, size: Tuple[int, int]):
        H, W = size
        dim = len(x.shape)
        if dim == 3: