Revert "Update __init__.py"

This reverts commit 2b699d44

megatron/__init__.py

@@ -26,7 +26,10 @@ from .initialize  import initialize_megatron
 
 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():
@@ -35,5 +38,8 @@ 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)