diff --git a/movinet.py b/movinet.py
index 3dbeb4e..793859f 100644
--- a/movinet.py
+++ b/movinet.py
@@ -106,7 +106,8 @@ class Movinet(NNOperator):
 
         self.model.clean_activation_buffers()
         feats = self.model.forward_features(inputs)
-        features = feats.to('cpu').squeeze(0).detach().numpy()
+        features = feats.flatten(1)
+        features = features.to('cpu').squeeze(0).detach().numpy()
 
         outs = self.model.head(feats)
         post_act = torch.nn.Softmax(dim=1)
diff --git a/result1.png b/result1.png
index 8730ba5..81320aa 100644
Binary files a/result1.png and b/result1.png differ
diff --git a/result2.png b/result2.png
index 6dfadb7..fb07273 100644
Binary files a/result2.png and b/result2.png differ