diff --git a/timm_image.py b/timm_image.py
index 286c9bc..1e869b9 100644
--- a/timm_image.py
+++ b/timm_image.py
@@ -144,15 +144,14 @@ class TimmImage(NNOperator):
         return img
 
     def post_proc(self, features):
-        features = features.to(self.device)
+        features = features.to('cpu')
         if features.dim() == 3:
             features = features[:, 0]
         if features.dim() == 4:
-            global_pool = nn.AdaptiveAvgPool2d(1).to(self.device)
+            global_pool = nn.AdaptiveAvgPool2d(1)
             features = global_pool(features)
             features = features.flatten(1)
         assert features.dim() == 2, f'Invalid output dim {features.dim()}'
-        features = features.to('cpu')
         return features
 
     def save_model(self, format: str = 'pytorch', path: str = 'default'):