fix: 修复部分已知问题

khaosz/trainer/train_callback.py

@@ -74,19 +74,16 @@ class SchedulerCallback(TrainCallback):
     Scheduler callback for trainer.
     """
     def __init__(self):
-        self.scheduler: LRScheduler = None
+        pass
     
     def on_train_begin(self, context: TrainContext):
         for group in context.optimizer.param_groups:
             if "initial_lr" not in group:
                 group["initial_lr"] = group["lr"] 
-        
-        self.scheduler = context.scheduler
     
     def on_batch_end(self, context: TrainContext):
-        _ = context
-        if self.scheduler:
-            self.scheduler.step()
+        if context.scheduler:
+            context.scheduler.step()
 
 
 class CheckpointCallback(TrainCallback):

khaosz/trainer/train_context.py

@@ -87,7 +87,7 @@ class TrainContextBuilder:
         return self
     
     def with_strategy(self) -> Self:
-        self._context.strategy = StrategyFactory.load(
+        self._context.strategy = StrategyFactory.create(
             model=self._context.model,
             train_type=self.config.strategy,
             device=get_current_device(),