fix(encoder): fix vald encocer

gnes-ai · Sep 9, 2019 · 814b2ee · 814b2ee
1 parent ffc822b
commit 814b2ee
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/gnes/encoder/numeric/vlad.py b/gnes/encoder/numeric/vlad.py
@@ -47,8 +47,8 @@ def faiss_index(self):
 
     def kmeans_pred(self, vecs):
         if self.using_faiss_pred:
-            D, I = self.index_flat.search(vecs.astype(np.float32), 1)
-            return np.reshape(I, [-1])
+            _, pred = self.index_flat.search(vecs.astype(np.float32), 1)
+            return np.reshape(pred, [-1])
         else:
             vecs = np.reshape(vecs, [vecs.shape[0], 1, 1, vecs.shape[1]])
             dist = np.sum(np.square(vecs - self.centroids), -1)