From bfbf9fdd9fe879466bebad8421c48b24e670d83b Mon Sep 17 00:00:00 2001
From: ChairC <974833488@qq.com>
Date: Mon, 18 Dec 2023 21:50:42 +0800
Subject: [PATCH] Update: The utils.initializer encapsulation method is used in
 distributed training.

---
 tools/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tools/train.py b/tools/train.py
index bcfdafa..7ee54e1 100644
--- a/tools/train.py
+++ b/tools/train.py
@@ -78,7 +78,7 @@ def train(rank=None, args=None):
         dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo", rank=rank,
                                 world_size=world_size)
         # Set device ID
-        device = torch.device("cuda", rank)
+        device = device_initializer(device_id=rank)
         # There may be random errors, using this function can reduce random errors in cudnn
         # torch.backends.cudnn.deterministic = True
         # Synchronization during distributed training