diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py
index f413ea7..c5878be 100644
--- a/convnext/models/convnext_tf.py
+++ b/convnext/models/convnext_tf.py
@@ -140,6 +140,7 @@ class ConvNeXt(tf.keras.Model):
 
     def forward_features(self, x):
         for i in range(len(self.downsample_layers)):
+            # x = self.downsample_layers[i](x)
             x = self.downsample_layers[i][0](x)
             x = self.downsample_layers[i][1](x)
             x = self.stages[i](x)
@@ -147,10 +148,6 @@ class ConvNeXt(tf.keras.Model):
 
     def call(self, x):
         x = self.forward_features(x)
-        if self.include_top:
-            x = self.avg(x)
-            x = self.norm(x)
-            x = self.head(x)
         return x