feat(trainer): 增加state_dict 存储设定

2026-02-04 19:47:21 +08:00
parent 7a9b9d0659
commit a5869d89ba
3 changed files with 13 additions and 4 deletions
@@ -105,6 +105,11 @@ class TrainConfig:
        default=None,
        metadata={"help": "Parallel function for training."}
    )
+    state_dict_wrapper: Optional[Callable] = field(
+         default=None,
+         metadata={"help": "Parallel function for state  dict saving."}
+    )
+    
    optimizer_factory: Optional[Callable[[nn.Module], Optimizer]] = field(
        default=None,
        metadata={"help": "Optimizer factory for training."}