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

---
 timm/models/davit.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index d607f945..0e2ea4cc 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -456,6 +456,13 @@ class DaViT(nn.Module):
         self.norms = norm_layer(self.num_features)
         self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=drop_rate)
         self.apply(self._init_weights)
+        
+        
+        
+        if self.features_only == True:
+            self.forward = self.forward_features_full
+        else:
+            self.forward = self.forward_classification
 
     def _init_weights(self, m):
         if isinstance(m, nn.Linear):
@@ -574,12 +581,13 @@ class DaViT(nn.Module):
         x = self.forward_features(x)
         x = self.forward_head(x)
         return x
-    
+    '''
     def forward(self, x):
         if self.features_only == True:
             return self.forward_features_full(x)
         else:
             return self.forward_classification(x)
+    '''
 
 def checkpoint_filter_fn(state_dict, model):
     """ Remap MSFT checkpoints -> timm """