diff --git a/omnivore.py b/omnivore.py
index 87b79ed..4759a75 100644
--- a/omnivore.py
+++ b/omnivore.py
@@ -101,10 +101,10 @@ class Omnivore(NNOperator):
                     )
         inputs = data.to(self.device)[None, ...]
 
-        feats = self.model.forward_features(inputs ,input_type = self.input_type)
+        feats = self.model.forward_features(inputs, input_type = self.input_type)
         features = feats.to('cpu').squeeze(0).detach().numpy()
 
-        outs = self.model.head(feats)
+        outs = self.model.head(feats, input_type = self.input_type)
         post_act = torch.nn.Softmax(dim=1)
         preds = post_act(outs)
         pred_scores, pred_classes = preds.topk(k=self.topk)