add omnivore

Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
2 years ago · 847b33d5fb
1 changed files with 1 additions and 1 deletions
--- a/omnivore.py
+++ b/omnivore.py
@ -99,7 +99,7 @@ class Omnivore(NNOperator):
                    )
        inputs = data.to(self.device)[None, ...]

-        outs = self.model.head(inputs)
+        outs = self.model(inputs)
        post_act = torch.nn.Softmax(dim=1)
        preds = post_act(outs)
        pred_scores, pred_classes = preds.topk(k=self.topk)