1 files changed, 1 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index 3776eb2..19348e5 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -535,6 +535,7 @@ def main():
     ]
 
     new_tokens = tokenizer.add_multi_tokens(args.placeholder_token, args.num_vectors)
+    embeddings.resize(len(tokenizer))
 
     for (new_token, init_ids) in zip(new_tokens, initializer_token_ids):
         embeddings.add_embed(new_token.ids, init_ids)