diff --git a/generate.py b/generate.py
index af46fd55..cc3ed0f8 100644
--- a/generate.py
+++ b/generate.py
@@ -19,7 +19,10 @@ def generate(tokenizer, prompt, model, config):
 def setup_model(config):
     model = AutoModelForCausalLM.from_pretrained(config["model_name"], device_map="auto", torch_dtype=torch.float16)
     tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"])
-    tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>"})
+    added_tokens = tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>", "pad_token": "<pad>"})
+
+    if added_tokens > 0:
+        model.resize_token_embeddings(len(tokenizer))
 
     if config["lora"]:
         model = PeftModelForCausalLM.from_pretrained(model, config["lora_path"], device_map="auto", torch_dtype=torch.float16)