From a02401aa758b9eecfb58ff3d6c5e0a0ac8a39b18 Mon Sep 17 00:00:00 2001 From: shuiyu1123 Date: Thu, 18 Jun 2026 18:40:25 +0000 Subject: [PATCH] fix: LLMTrainer supports lr param --- llm_chat_improved.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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) # 混合精度训练