tensorizer return parameter fix
@@ -777,7 +777,8 @@ class TextCategorizer(Pipe):
def predict(self, docs):
scores = self.model(docs)
scores = self.model.ops.asarray(scores)
- return scores
+ tensors = [doc.tensor for doc in docs]
+ return scores, tensors
def set_annotations(self, docs, scores, tensors=None):
for i, doc in enumerate(docs):
