Addressed Jared's comments

tasks/zeroshot_gpt2/datasets.py

@@ -21,7 +21,8 @@ import math
 import numpy as np
 import torch
 
-from megatron import get_args, print_rank_0
+from megatron import get_args
+from megatron import print_rank_0
 from megatron import get_tokenizer
 from .detokenizer import get_detokenizer
 

tasks/zeroshot_gpt2/evaluate.py

@@ -19,7 +19,8 @@ import math
 
 import torch
 
-from megatron import get_args, print_rank_0
+from megatron import get_args
+from megatron import print_rank_0
 from megatron import get_tokenizer
 from megatron import mpu
 from megatron.checkpointing import load_checkpoint

tools/generate_samples_gpt2.py

@@ -20,7 +20,8 @@ import sys
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                              os.path.pardir)))
 
-from megatron import get_args, print_rank_0
+from megatron import get_args
+from megatron import print_rank_0
 from megatron import get_tokenizer
 from megatron.checkpointing import load_checkpoint
 from megatron.initialize import initialize_megatron