Revert "Update __init__.py"

def print_rank_0(message):

    """If distributed is initialized, print only on rank 0."""

    if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:

    if torch.distributed.is_initialized():

        if torch.distributed.get_rank() == 0:

            print(message, flush=True)

    else:

        print(message, flush=True)

def is_last_rank():

def print_rank_last(message):

    """If distributed is initialized, print only on last rank."""

    if torch.distributed.is_initialized() and is_last_rank():

    if torch.distributed.is_initialized():

        if is_last_rank():

            print(message, flush=True)

    else:

        print(message, flush=True)