diff --git a/omnivore.py b/omnivore.py index 1d2ab94..584c919 100644 --- a/omnivore.py +++ b/omnivore.py @@ -99,9 +99,9 @@ class Omnivore(NNOperator): ) inputs = data.to(self.device)[None, ...] - outs = self.model(inputs) + outs = self.model(inputs,input_type="video") post_act = torch.nn.Softmax(dim=1) - preds = post_act(outs,input_type="video") + preds = post_act(outs) pred_scores, pred_classes = preds.topk(k=self.topk) labels = [self.classmap[int(i)] for i in pred_classes[0]] scores = [round(float(x), 5) for x in pred_scores[0]]