From c48dd0e1d33416915d86d2728bd2d9545b1940b1 Mon Sep 17 00:00:00 2001
From: Matthew Honnibal <honnibal+gh@gmail.com>
Date: Wed, 1 Nov 2017 02:06:58 +0100
Subject: [PATCH] Fix vector pruning

---
 spacy/vectors.pyx | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/spacy/vectors.pyx b/spacy/vectors.pyx
index 95378947a..a77fb2236 100644
--- a/spacy/vectors.pyx
+++ b/spacy/vectors.pyx
@@ -275,7 +275,10 @@ cdef class Vectors:
             sims = xp.dot(batch, vectors.T)
             best_rows[i:i+batch_size] = sims.argmax(axis=1)
             scores[i:i+batch_size] = sims.max(axis=1)
-        keys = self.find(rows=best_rows)
+
+        xp = get_array_module(self.data)
+        row2key = {row: key for key, row in self.key2row.items()}
+        keys = xp.asarray([row2key[row] for row in best_rows], dtype='uint64')
         return (keys, best_rows, scores)
 
     def from_glove(self, path):