修复多机训练问题 (#143)

* 修复多机训练问题

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

train_ms.py

@@ -62,6 +62,7 @@ def run():
         init_method="env://",  # If Windows,switch to gloo backend.
     )  # Use torchrun instead of mp.spawn
     rank = dist.get_rank()
+    local_rank = int(os.environ["LOCAL_RANK"])
     n_gpus = dist.get_world_size()
 
     # 命令行/config.yml配置解析
@@ -99,7 +100,8 @@ def run():
             f.write(data)
 
     torch.manual_seed(hps.train.seed)
-    torch.cuda.set_device(rank)
+    torch.cuda.set_device(local_rank)
+
     global global_step
     if rank == 0:
         logger = utils.get_logger(hps.model_dir)