fix: 修复参数传递问题并更新测试单元

2026-02-28 19:01:16 +08:00
parent b17cc6a6fb
commit 6089a12cef
4 changed files with 25 additions and 21 deletions
@@ -47,7 +47,7 @@ class TrainContextBuilder:
            fn = self.config.parallel_wrapper
            self._context.model = fn(self._context.model)
        
-        self._context.optimizer = self.config.optimizer_fn(self._context.model.parameters())
+        self._context.optimizer = self.config.optimizer_fn(self._context.model)
        self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
    
    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self: