[DLMED] fix log print (#212)

Signed-off-by: Nic Ma <nma@nvidia.com>

Author: Nic-Ma

acceleration/distributed_training/unet_evaluation_workflows.py

@@ -82,6 +82,7 @@
 
 
 def evaluate(args):
+    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
     if args.local_rank == 0 and not os.path.exists(args.dir):
         # create 16 random image, mask paris for evaluation
         print(f"generating synthetic data to {args.dir} (this may take a while)")
@@ -149,7 +150,6 @@ def evaluate(args):
         ),
     ]
     if dist.get_rank() == 0:
-        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
         val_handlers.extend(
             [
                 StatsHandler(output_transform=lambda x: None),

acceleration/distributed_training/unet_training_workflows.py

@@ -84,6 +84,7 @@
 
 
 def train(args):
+    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
     if args.local_rank == 0 and not os.path.exists(args.dir):
         # create 40 random image, mask paris for training
         print(f"generating synthetic data to {args.dir} (this may take a while)")
@@ -160,7 +161,6 @@ def train(args):
         LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
     ]
     if dist.get_rank() == 0:
-        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
         train_handlers.extend(
             [
                 StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]),