Function print_rank0

SwissArmyTransformer/sat/helpers.py:127–137

(msg, level=logging.INFO, flush=True)

logger = configure_logging()

def print_rank0(msg, level=logging.INFO, flush=True):
    if isinstance(level, str):
        level = getattr(logging, level.upper())
    if torch.distributed.is_initialized():
        msg = f"[RANK {torch.distributed.get_rank()}] {msg}"
        if torch.distributed.get_rank() == 0:
            logger.log(level=level, msg=msg)
            if flush:
                logger.handlers[0].flush()
    else:
        logger.log(level=level, msg=msg)

def print_all(msg, level=logging.INFO, flush=True):
    if isinstance(level, str):