code: def set_logger(args):
      """
      Write logs to checkpoint and console
      """
  
      if args.do_train:
          log_file = os.path.join(
              args.save_path or args.init_checkpoint, "train.log"
          )
      else:
          log_file = os.path.join(
              args.save_path or args.init_checkpoint, "test.log"
          )
  
      logging.basicConfig(
          format="%(asctime)s %(levelname)-8s %(message)s",
          level=logging.INFO,
          datefmt="%Y-%m-%d %H:%M:%S",
          filename=log_file,
          filemode="w",
      )
  
      if args.print_on_screen:
          console = logging.StreamHandler()
          console.setLevel(logging.INFO)
          formatter = logging.Formatter(
              "%(asctime)s %(levelname)-8s %(message)s"
          )
          console.setFormatter(formatter)
          logging.getLogger("").addHandler(console)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import argparse
  import json
  import logging
  import os
  import random
  import numpy
  import torch
  from torch.utils.data import DataLoader
  from model import KGEModel
  from dataloader import TrainDataset
  from dataloader import BidirectionalOneShotIterator
  from ogb.linkproppred import LinkPropPredDataset, Evaluator
  from collections import defaultdict
  from tqdm.auto import tqdm
  import time
  from tensorboardX import SummaryWriter
label_desc: Use image node_id 4 for calling a global function with example usage: set_logger(args) without return types