fix dtype mismatch error

CPJKU · Jan 21, 2022 · c985515 · c985515
1 parent d9f0752
commit c985515
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/wechsel/__init__.py b/wechsel/__init__.py
@@ -240,7 +240,9 @@ def get_n_closest(token_id, similarities, top_k):
 
     source_vocab = source_tokenizer.vocab
 
-    target_matrix = np.zeros((len(target_tokenizer), source_matrix.shape[1]))
+    target_matrix = np.zeros(
+        (len(target_tokenizer), source_matrix.shape[1]), dtype=source_matrix.dtype
+    )
 
     mean, std = (
         source_matrix.mean(0),