Skip to content

GitLab

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
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 8 deletions
megatron/__init__.py
View file @ 2b699d44
...@@ -26,10 +26,7 @@ from .initialize import initialize_megatron ...@@ -26,10 +26,7 @@ from .initialize import initialize_megatron
def print_rank_0(message): def print_rank_0(message):
"""If distributed is initialized, print only on rank 0.""" """If distributed is initialized, print only on rank 0."""
if torch.distributed.is_initialized(): if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
if torch.distributed.get_rank() == 0:
print(message, flush=True)
else:
print(message, flush=True) print(message, flush=True)
def is_last_rank(): def is_last_rank():
...@@ -38,8 +35,5 @@ def is_last_rank(): ...@@ -38,8 +35,5 @@ def is_last_rank():
def print_rank_last(message): def print_rank_last(message):
"""If distributed is initialized, print only on last rank.""" """If distributed is initialized, print only on last rank."""
if torch.distributed.is_initialized(): if torch.distributed.is_initialized() and is_last_rank():
if is_last_rank():
print(message, flush=True)
else:
print(message, flush=True) print(message, flush=True)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment