From 15b168e3055c7ea42e839a2a192ee35278491bbe Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Thu, 8 Dec 2022 01:32:06 -0800
Subject: [PATCH] Update davit.py

---
 timm/models/davit.py | 34 ++++++++++++++++------------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 37870b1e..e31b5e58 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -474,7 +474,21 @@ class DaViT(nn.Module):
             return self.forward_classification
     '''
     
-
+    @torch.jit.ignore
+    def _get_forward_fn(self):
+        if self._features_only == True:
+            return self.forward_features_full
+        else:
+            return self.forward_classification
+    
+    @property
+    def features_only(self):
+        return self._features_only
+        
+    @features_only.setter
+    def features_only(self, new_value : bool):
+        self._features_only = new_value
+        self.forward = self._get_forward_fn()
     
     
     
@@ -577,7 +591,7 @@ class DaViT(nn.Module):
         
         
         # non-normalized pyramid features + corresponding sizes
-        return features, sizes
+        return features[1:], sizes[:-1]
         
     def forward_features(self, x):
         x, sizes = self.forward_features_full(x)
@@ -598,22 +612,6 @@ class DaViT(nn.Module):
     
     def forward(self, x):
         return x
-        
-    @torch.jit.ignore
-    def _get_forward_fn(self):
-        if self._features_only == True:
-            return self.forward_features_full
-        else:
-            return self.forward_classification
-    
-    @property
-    def features_only(self):
-        return self._features_only
-        
-    @features_only.setter
-    def features_only(self, new_value : bool):
-        self._features_only = new_value
-        self.forward = self._get_forward_fn()
     
 
 def checkpoint_filter_fn(state_dict, model):