1 files changed, 5 insertions, 1 deletions

diff --git a/training/functional.py b/training/functional.py
index 27a43c2..4565612 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -231,12 +231,16 @@ def add_placeholder_tokens(
     embeddings: ManagedCLIPTextEmbeddings,
     placeholder_tokens: list[str],
     initializer_tokens: list[str],
-    num_vectors: Union[list[int], int]
+    num_vectors: Optional[Union[list[int], int]] = None,
 ):
     initializer_token_ids = [
         tokenizer.encode(token, add_special_tokens=False)
         for token in initializer_tokens
     ]
+
+    if num_vectors is None:
+        num_vectors = [len(ids) for ids in initializer_token_ids]
+
     placeholder_token_ids = tokenizer.add_multi_tokens(placeholder_tokens, num_vectors)
 
     embeddings.resize(len(tokenizer))