diff --git a/llm_chat_improved.py b/llm_chat_improved.py
index 89521fe..9e30385 100644
--- a/llm_chat_improved.py
+++ b/llm_chat_improved.py
@@ -580,11 +580,11 @@ class ConversationDataset(Dataset):
 class LLMTrainer:
     """LLM训练器"""
     
-    def __init__(self, model: nn.Module, device: str = 'cpu', use_amp: bool = None):
+    def __init__(self, model: nn.Module, device: str = 'cpu', lr: float = 0.0001, use_amp: bool = None):
         self.model = model.to(device)
         self.device = device
         self.criterion = nn.CrossEntropyLoss(ignore_index=0)
-        self.optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
+        self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
         self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
         
         # 混合精度训练