From 847b33d5fb37384d56bae0bdbbe3a65b21a72865 Mon Sep 17 00:00:00 2001 From: gexy5 Date: Tue, 14 Jun 2022 15:05:56 +0800 Subject: [PATCH] add omnivore Signed-off-by: gexy5 --- omnivore.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/omnivore.py b/omnivore.py index 6aa49ec..fac3b53 100644 --- 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)