Browse Source

[misc] fix dist logger (#5782)

pull/5787/head
Hongxin Liu 6 months ago committed by GitHub
parent
commit
b9d646fe9e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 2
      colossalai/nn/optimizer/__init__.py

2
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):

Loading…
Cancel
Save