Update davit.py

timm/models/davit.py

@@ -103,7 +103,7 @@ class PatchEmbed(nn.Module):
     @register_notrace_function # reason: dim in control sequence
     def forward(self, x : Tensor, size: Tuple[int, int]):
         H, W = size
-        dim = len(x.shape)
+        dim = x.dim()
         if dim == 3:
             B, HW, C = x.shape
             x = self.norm(x)