From d53a8d9313099b0c9724e28ca276603274749313 Mon Sep 17 00:00:00 2001 From: adrianeboyd Date: Mon, 7 Oct 2019 13:38:35 +0200 Subject: [PATCH] Consider batch_size when sorting similar vectors (#4388) --- spacy/vectors.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spacy/vectors.pyx b/spacy/vectors.pyx index 75716617c..881f01052 100644 --- a/spacy/vectors.pyx +++ b/spacy/vectors.pyx @@ -337,7 +337,7 @@ cdef class Vectors: scores[i:i+batch_size] = xp.partition(sims, -n, axis=1)[:,-n:] if sort: - sorted_index = xp.arange(scores.shape[0])[:,None],xp.argsort(scores, axis=1)[:,::-1] + sorted_index = xp.arange(scores.shape[0])[:,None],xp.argsort(scores[i:i+batch_size], axis=1)[:,::-1] scores[i:i+batch_size] = scores[sorted_index] best_rows[i:i+batch_size] = best_rows[sorted_index]