Unverified Commit 2b699d44 authored by Satpal Singh Rathore's avatar Satpal Singh Rathore Committed by GitHub
Browse files

Update __init__.py 

Removed redundant if-else.
parent e64507f0

megatron/__init__.py

+2 −8
Original line number Diff line number Diff line
@@ -26,10 +26,7 @@ from .initialize import initialize_megatron 


def print_rank_0(message):

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

    if torch.distributed.is_initialized():

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

            print(message, flush=True)

    else:

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

        print(message, flush=True)



def is_last_rank():
@@ -38,8 +35,5 @@ def is_last_rank(): 


def print_rank_last(message):

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

    if torch.distributed.is_initialized():

        if is_last_rank():

            print(message, flush=True)

    else:

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

        print(message, flush=True)