fix: 修复参数传递问题

khaosz/parallel/setup.py

@@ -81,6 +81,7 @@ def only_on_rank(rank, sync=False):
     def decorator(func):
         @wraps(func)
         def wrapper(*args, **kwargs):
+            ret_args = None
             if get_rank() == rank:
                 ret_args = func(*args, **kwargs)
 

khaosz/trainer/trainer.py

@@ -57,6 +57,8 @@ class Trainer:
             world_size=config.nprocs,
             master_addr=config.master_addr,
             master_port=config.master_port,
+            device_type=config.device_type,
+            device_ids=config.device_ids,
             checkpoint=checkpoint
         )