refactor(paralell): 优化并行设备指定方法

khaosz/trainer/train_context.py

@@ -88,13 +88,13 @@ class TrainContextBuilder:
         )
         return self
     
-    def with_parallel_fn(self) -> Self:
+    def with_parallel(self) -> Self:
         device = get_current_device()
         self._context.model = self._context.model.to(device=device)
         
         if self.config.nprocs > 1:
             
-            fn = self.config.parallel_fn
+            fn = self.config.parallel_wrapper
             optimizer_fn = self.config.optimizer_factory
             scheduler_fn = self.config.scheduler_factory
             

khaosz/trainer/trainer.py

@@ -38,7 +38,7 @@ class Trainer:
                 .with_checkpoint(checkpoint)
                 .with_dataloader()
                 .with_strategy()
-                .with_parallel_fn()
+                .with_parallel()
                 .build())
     
     def _call_callbacks(self, method_name: str, context: TrainContext):