From 2a66366be4c2715870e4859fd5a5db6e8a9dc00a Mon Sep 17 00:00:00 2001
From: chenmengzheAAA <123789350+chenmengzheAAA@users.noreply.github.com>
Date: 星期四, 14 九月 2023 19:00:17 +0800
Subject: [PATCH] Merge pull request #956 from alibaba-damo-academy/chenmengzheAAA-patch-4

---
 funasr/train/distributed_utils.py |    8 ++------
 1 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/funasr/train/distributed_utils.py b/funasr/train/distributed_utils.py
index 088203a..13f5744 100644
--- a/funasr/train/distributed_utils.py
+++ b/funasr/train/distributed_utils.py
@@ -53,15 +53,13 @@
             # https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group
             os.environ.setdefault("NCCL_BLOCKING_WAIT", "1")
 
-            torch.distributed.init_process_group(backend='nccl',
+            torch.distributed.init_process_group(backend=self.dist_backend,
                                                  init_method=self.dist_init_method,
                                                  world_size=args.dist_world_size,
                                                  rank=args.dist_rank)
             self.dist_rank = torch.distributed.get_rank()
             self.dist_world_size = torch.distributed.get_world_size()
             self.local_rank = args.local_rank
-            logging.info("world size: {}, rank: {}, local_rank: {}".format(self.dist_world_size, self.dist_rank,
-                                                                           self.local_rank))
 
     def init_options_pai(self):
         if self.distributed:
@@ -115,12 +113,10 @@
             # https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group
             os.environ.setdefault("NCCL_BLOCKING_WAIT", "1")
 
-            torch.distributed.init_process_group(backend='nccl', init_method='env://')
+            torch.distributed.init_process_group(backend=self.dist_backend, init_method='env://')
             self.dist_rank = torch.distributed.get_rank()
             self.dist_world_size = torch.distributed.get_world_size()
             self.local_rank = args.local_rank
-            logging.info("world size: {}, rank: {}, local_rank: {}".format(self.dist_world_size, self.dist_rank,
-                                                                           self.local_rank))
 
 
 def resolve_distributed_mode(args):

--
Gitblit v1.9.1