Consider batch_size when sorting similar vectors
@@ -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]
