diff --git a/timm_image.py b/timm_image.py
index 94d1eec..b9e03cd 100644
--- a/timm_image.py
+++ b/timm_image.py
@@ -128,7 +128,7 @@ class TimmImage(NNOperator):
         return img
 
     def post_proc(self, features):
-        if 'vit' in self.model_name and features.dim() == 3:
+        if features.dim() == 3:
             features = features[:, 0]
         if features.dim() == 4:
             global_pool = nn.AdaptiveAvgPool2d(1).to(self.device)