Update davit.py

3 years ago · 1d2a34f20f
parent 19a87cae63
commit 1d2a34f20f
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -100,7 +100,7 @@ class PatchEmbed(nn.Module):
                padding=to_2tuple(pad))
            self.norm = nn.LayerNorm(in_chans)

-    @register_notrace_function # reason: dim in control sequence
+@register_notrace_function # reason: dim in control sequence
    def forward(self, x : Tensor, size: Tuple[int, int]):
        H, W = size
        dim = len(x.shape)