Skip to content

GitLab

Unverified Commit 8f241a96 authored by Satpal Singh Rathore's avatar Satpal Singh Rathore Committed by GitHub
Browse files

Revert "Update __init__.py" 

This reverts commit 2b699d44
parent d03059a6
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 2 deletions
megatron/__init__.py
View file @ 8f241a96
...@@ -26,7 +26,10 @@ from .initialize import initialize_megatron ...@@ -26,7 +26,10 @@ 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() 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) print(message, flush=True)
def is_last_rank(): def is_last_rank():
...@@ -35,5 +38,8 @@ def is_last_rank(): ...@@ -35,5 +38,8 @@ 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() and is_last_rank(): if torch.distributed.is_initialized():
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