code: def train(args, dataset, generator, discriminator):
      step = int(math.log2(args.init_size)) - 2
      resolution = 4 * 2**step
      loader = sample_data(
          dataset,
          args.batch.get(resolution, args.batch_default),
          resolution,
      )
      data_loader = iter(loader)
  
      adjust_lr(g_optimizer, args.lr.get(resolution, args.lr_default))
      adjust_lr(d_optimizer, args.lr.get(resolution, args.lr_default))
  
      alpha = 0
      used_sample = 0
  
      max_step = int(math.log2(args.max_size)) - 2
      final_progress = False
  
      requires_grad(generator, False)
      requires_grad(discriminator, True)
  
      pbar = tqdm(range(200_000))
  
      for i in pbar:
          alpha = min(1, 1 / args.phase * (used_sample + 1))
  
          if (
              resolution == args.init_size and args.ckpt is None
          ) or final_progress:
              alpha = 1
  
          if used_sample > args.phase * 2:
              used_sample = 0
              step += 1
  
              if step > max_step:
                  step = max_step
                  final_progress = True
                  ckpt_step = step + 1
  
              else:
                  alpha = 0
                  ckpt_step = step
  
              resolution = 4 * 2**step
  
              loader = sample_data(
                  dataset,
                  args.batch.get(resolution, args.batch_default),
                  resolution,
              )
              data_loader = iter(loader)
  
              generator.save_parameters(
                  osp.join(
                      args.ckpt_dir,
                      f"generator_step-{ckpt_step}.params",
                  )
              )
              discriminator.save_parameters(
                  osp.join(
                      args.ckpt_dir,
                      f"discriminator_step-{ckpt_step}.params",
                  )
              )
              g_running.save_parameters(
                  osp.join(
                      args.ckpt_dir,
                      f"g_running_step-{ckpt_step}.params",
                  )
              )
  
              adjust_lr(
                  g_optimizer, args.lr.get(resolution, args.lr_default)
              )
              adjust_lr(
                  d_optimizer, args.lr.get(resolution, args.lr_default)
              )
  
          try:
              real_image = next(data_loader)
  
          except (OSError, StopIteration):
              data_loader = iter(loader)
              real_image = next(data_loader)
  
          used_sample += real_image.shape[0]
          b_size = real_image.shape[0]
          real_image_list = gluon.utils.split_and_load(
              real_image, ctx_list=context, batch_axis=0
          )
  
          if args.mixing and random.random() < 0.9:
              gen_in11, gen_in12, gen_in21, gen_in22 = nd.random.randn(
                  4, b_size, code_size
              ).split(4, 0)
              gen_in1 = [
                  nd.squeeze(gen_in11, axis=0),
                  nd.squeeze(gen_in12, axis=0),
              ]
              gen_in2 = [
                  nd.squeeze(gen_in21, axis=0),
                  nd.squeeze(gen_in22, axis=0),
              ]
  
          else:
              gen_in1, gen_in2 = nd.random.randn(
                  2, b_size, code_size
              ).split(2, 0)
              gen_in1 = nd.squeeze(gen_in1, axis=0)
              gen_in2 = nd.squeeze(gen_in2, axis=0)
  
          gen_in1_list = gluon.utils.split_and_load(
              gen_in1, ctx_list=context, batch_axis=0
          )
          gen_in2_list = gluon.utils.split_and_load(
              gen_in2, ctx_list=context, batch_axis=0
          )
  
          if args.loss == "wgan":
              fake_predict_list = []
              real_predict_list = []
              D_loss_list = []
              with autograd.record():
                  for _, (rl_image, g1) in enumerate(
                      zip(real_image_list, gen_in1_list)
                  ):
                      real_predict = discriminator(
                          rl_image, step, alpha
                      )
                      real_predict = -real_predict.mean()
                      real_predict_list.append(real_predict)
  
                      fake_image = generator(g1, step, alpha)
                      fake_predict = discriminator(
                          fake_image.detach(), step, alpha
                      )
                      fake_predict = fake_predict.mean()
                      fake_predict_list.append(fake_predict)
  
                      D_loss_list.append(real_predict + fake_predict)
  
              autograd.backward(loss_list)
  
          elif args.loss == "r1":
              # Not able to implement r1 loss
              raise Exception(
                  "r1 loss has not been implemented, please use wgan loss"
              )
          else:
              raise Exception("Not valid loss, please use wgan loss")
  
          if i % 10 == 0:
              real_predict_val = [
                  i.asnumpy() for i in real_predict_list
              ]
              fake_predict_val = [
                  i.asnumpy() for i in fake_predict_list
              ]
              d_real_val = np.concatenate(real_predict_val).mean()
              d_fake_val = np.concatenate(fake_predict_val).mean()
              disc_loss_val = d_real_val + d_fake_val
  
          d_optimizer.step(b_size, ignore_stale_grad=True)
  
          if (i + 1) % n_critic == 0:
              requires_grad(generator, True)
              requires_grad(discriminator, False)
  
              if args.loss == "wgan-gp":
                  predict_list = []
                  with autograd.record():
                      for _, g2 in enumerate(gen_in2_list):
                          fake_image = generator(g2, step, alpha)
                          predict = discriminator(
                              fake_image, step, alpha
                          )
                          predict = -predict.mean()
                          predict_list.append(predict)
                  autograd.backward(predict_list)
              elif args.loss == "r1":
                  # Not able to implement r1 loss
                  raise Exception(
                      "r1 loss has not been implemented, please use wgan loss"
                  )
              else:
                  raise Exception(
                      "Not valid loss, please use wgan loss"
                  )
  
              if i % 10 == 0:
                  predict_val = [i.asnumpy() for i in predict_list]
                  gen_loss_val = np.concatenate(predict_val).mean()
  
              g_optimizer.step(b_size, ignore_stale_grad=True)
  
              accumulate(g_running, generator)
  
              requires_grad(generator, False)
              requires_grad(discriminator, True)
  
          if (i + 1) % 100 == 0:
              images = []
  
              gen_i, gen_j = args.gen_sample.get(resolution, (10, 5))
  
              for _ in range(gen_i):
                  results = g_running(
                      nd.random.randn(gen_j, code_size, ctx=mx.gpu(0)),
                      step,
                      alpha,
                  )
                  for r in results:
                      images.append(r)
  
              plot_images(
                  images,
                  osp.join(args.out, f"{str(i + 1).zfill(6)}.png"),
                  gen_i,
                  gen_j,
              )
  
          if (i + 1) % 1000 == 0:
              generator.save_parameters(
                  osp.join(
                      args.ckpt_dir, f"g-{str(i + 1).zfill(6)}.params"
                  )
              )
              discriminator.save_parameters(
                  osp.join(
                      args.ckpt_dir, f"d-{str(i + 1).zfill(6)}.params"
                  )
              )
              g_running.save_parameters(
                  osp.join(
                      args.ckpt_dir,
                      f"g_running-{str(i + 1).zfill(6)}.params",
                  )
              )
  
          state_msg = (
              f"Size: {4 * 2 ** step}; G: {gen_loss_val:.1f}; D: {disc_loss_val:.1f};"
              f"D_real: {d_real_val:.1f}; D_fake: {d_fake_val:.1f}; Alpha: {alpha:.4f}"
          )
  
          logger.info(
              f"Size: {4 * 2 ** step}; G: {gen_loss_val:.1f}; D: {disc_loss_val:.1f}\
              D_real: {d_real_val:1f}; D_fake: {d_fake_val:1f}; Alpha: {alpha:.4f}"
          )
  
          pbar.set_description(state_msg)
  
imports: import os
  import random
  import math
  import argparse
  import logging
  import os.path
  from io import BytesIO
  import lmdb
  import numpy
  from PIL import Image
  from tqdm import tqdm
  from mpl_toolkits.axes_grid1 import ImageGrid
  import matplotlib.pyplot
  import mxnet
  import mxnet.ndarray
  from mxnet import gluon, autograd
  from mxnet.gluon.data.vision import transforms
  from mxnet.gluon.data import DataLoader
  from model import StyledGenerator, Discriminator
label_desc: Use image node_id 7 for calling a global function with example usage: train(args, dataset, generator, discriminator) without return types