From 488c263a6301c135c1908e8299b17ead0183bb3e Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Thu, 8 Dec 2022 21:51:05 -0800
Subject: [PATCH] Update davit.py

---
 timm/models/davit.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 6ec0aed4..4f1f0914 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -100,10 +100,10 @@ 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)
+        dim = x.dim()
         if dim == 3:
             B, HW, C = x.shape
             x = self.norm(x)