From b9d646fe9e2880203cda6fa9b6465aef131db354 Mon Sep 17 00:00:00 2001
From: Hongxin Liu <lhx0217@gmail.com>
Date: Wed, 5 Jun 2024 15:04:22 +0800
Subject: [PATCH] [misc] fix dist logger (#5782)

---
 colossalai/nn/optimizer/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/colossalai/nn/optimizer/__init__.py b/colossalai/nn/optimizer/__init__.py
index ddb03f947..d872dbbaf 100644
--- a/colossalai/nn/optimizer/__init__.py
+++ b/colossalai/nn/optimizer/__init__.py
@@ -43,11 +43,11 @@ optim2DistOptim = {
     CAME: DistributedCAME,
     Adafactor: DistributedAdaFactor,
 }
-_logger = get_dist_logger()
 
 
 def cast_to_distributed(optim):
     if optim.__class__ in optim2DistOptim:
+        _logger = get_dist_logger()
         _logger.info(f"Converting optimizer {optim.__class__.__name__} to its distributed version.", ranks=[0])
 
         if isinstance(optim, GaLoreAdamW8bit):