From 95a483534258fc2d93de74296f7b172d5a6b87df Mon Sep 17 00:00:00 2001 From: Richard Hudson Date: Thu, 22 Dec 2022 15:36:14 +0100 Subject: [PATCH] Convert argmax result to raw integer Co-authored-by: Madeesh Kannan --- spacy/pipeline/edit_tree_lemmatizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spacy/pipeline/edit_tree_lemmatizer.py b/spacy/pipeline/edit_tree_lemmatizer.py index 8a26e4aba..45ed59420 100644 --- a/spacy/pipeline/edit_tree_lemmatizer.py +++ b/spacy/pipeline/edit_tree_lemmatizer.py @@ -165,7 +165,7 @@ class EditTreeLemmatizer(TrainablePipe): doc_compat_guesses = [] for i, token in enumerate(doc): for _ in range(self.top_k): - candidate = doc_scores[i].argmax() + candidate = int(doc_scores[i].argmax()) candidate_tree_id = self.cfg["labels"][candidate] if self.trees.apply(candidate_tree_id, token.text) is not None: doc_compat_guesses.append(candidate_tree_id)