fixed not being added while training

soldni · Oct 11, 2022 · 094f2dd · 094f2dd
1 parent 288d729
commit 094f2dd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/pyterrier_sentence_transformers/index.py b/src/pyterrier_sentence_transformers/index.py
@@ -67,10 +67,11 @@ def __init__(
     def index_data(self, ids: List[Any], embeddings: np.ndarray):
         self._update_id_mapping(ids)
         embeddings = embeddings.astype('float32')
+
         if not self.index.is_trained:
             self.index.train(embeddings)
-        else:
-            self.index.add(embeddings)
+
+        self.index.add(embeddings)
 
         print(f'Total data indexed {len(self.index_id_to_db_id)}')