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

---
 timm/models/davit.py | 32 +++++++++++++++++---------------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 15674187..37870b1e 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -474,21 +474,7 @@ 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()
+
     
     
     
@@ -612,6 +598,22 @@ 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):