logo
Browse Source

add omnivore

Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
main
gexy5 2 years ago
parent
commit
847b33d5fb
  1. 2
      omnivore.py

2
omnivore.py

@ -99,7 +99,7 @@ class Omnivore(NNOperator):
) )
inputs = data.to(self.device)[None, ...] inputs = data.to(self.device)[None, ...]
outs = self.model.head(inputs)
outs = self.model(inputs)
post_act = torch.nn.Softmax(dim=1) post_act = torch.nn.Softmax(dim=1)
preds = post_act(outs) preds = post_act(outs)
pred_scores, pred_classes = preds.topk(k=self.topk) pred_scores, pred_classes = preds.topk(k=self.topk)

Loading…
Cancel
Save