fix

Author: thiswillbeyourgithub

WDoc/utils/embeddings.py

@@ -92,7 +92,7 @@ def wrapper(vector, k):
             diff = k - new_ids.shape[0]
             if diff == 0:
                 assert original_scores.shape == original_ids.shape
-                assert original_ids.shape == new_ids.shape
+                assert original_ids.squeeze().shape == new_ids.shape
                 return original_scores, original_ids
 
             trial = 0