diff --git a/scripts/train_cnb.py b/scripts/train_cnb.py
index 5f13467..e4ca109 100644
--- a/scripts/train_cnb.py
+++ b/scripts/train_cnb.py
@@ -84,7 +84,7 @@ def main():
             delattr(tokenizer_instance, attr)
 
     # ====== 模型 ======
-    model = LLMModel(vocab_size=len(vocab), d_model=320, n_heads=8,
+    model = LLMModel(vocab_size=max(vocab.values()) + 1, d_model=320, n_heads=8,
                      n_layers=6, d_ff=1536, max_len=256, dropout=0.1,
                      use_memory=False, use_thinking=False, use_moe=False)
     total = sum(p.numel() for p in model.parameters())