From 2f48986bbe340112c3d0ebab1051e1f6d32a10eb Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Thu, 30 May 2024 16:16:19 +0800 Subject: [PATCH] Fix token count Signed-off-by: Jael Gu --- rerank.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rerank.py b/rerank.py index 85cb72f..7a44e36 100644 --- a/rerank.py +++ b/rerank.py @@ -67,7 +67,7 @@ class ReRank(NNOperator): texts[idx].append(text.strip()) tokenized = self.tokenizer(*texts, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_length) - token_count = torch.count_nonzero(tokenized['input_ids']) + token_count = int(torch.count_nonzero(tokenized['input_ids'])) logits = self.model(**tokenized) scores = self.post_proc(logits)