From 3e0477fd79690e253c64cc8e900d33959f179010 Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Sat, 6 Sep 2025 17:25:46 +0800
Subject: [PATCH] update

---
 train_pretrain_accelerate.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/train_pretrain_accelerate.py b/train_pretrain_accelerate.py
index 22600cd..117b8b3 100644
--- a/train_pretrain_accelerate.py
+++ b/train_pretrain_accelerate.py
@@ -1421,12 +1421,14 @@ def main():
     lm_config = LMConfig(
         dim=args.dim,
         n_layers=args.n_layers,
+        n_heads=args.n_heads,
         max_seq_len=args.max_seq_len,
         use_moe=args.use_moe,
         disable_db=args.disable_db,
         flash_attn=args.use_flash_attn,
         knowledge_num=args.knowledge_num,
         knowledge_length=args.knowledge_length,
+        knowledge_dim=args.knowledge_dim,
         embeddings_epoch=args.embedding_epoch,
         freeze_ratio=args.freeze_ratio
     )