code: def train(device, data_path, output_dir, num_layers):
      logger.info("Build dataloader")
      (
          train_dataset,
          valid_dataset,
          test_dataset,
          embedding,
      ) = read_data_sst(
          data_path,
          FLAGS.max_input_length,
          FLAGS.min_count,
          train_ratio=FLAGS.train_ratio,
          valid_ratio=FLAGS.valid_ratio,
          is_binary=FLAGS.is_binary,
      )
      train_dataloader = DataLoader(
          train_dataset,
          batch_size=FLAGS.batch_size,
          shuffle=True,
          pin_memory=True,
      )
      test_dataloader = DataLoader(
          test_dataset,
          batch_size=FLAGS.eval_batch_size,
          pin_memory=True,
      )
      valid_dataloader = DataLoader(
          valid_dataset,
          batch_size=FLAGS.eval_batch_size,
          pin_memory=True,
      )
  
      logger.info("Build model")
      child_model = get_model(embedding, num_layers)
      logger.info("Finish build model")
  
      # for name, var in child_model.named_parameters():
      #    logger.info(name, var.size(), var.requires_grad)  # output all params
  
      num_vars = count_model_params(child_model.parameters())
      logger.info("Model has {} params".format(num_vars))
  
      for m in child_model.modules():  # initializer
          if isinstance(m, (nn.Conv1d, nn.Linear)):
              nn.init.xavier_uniform_(m.weight)
  
      criterion = nn.CrossEntropyLoss()
  
      # get optimizer
      if FLAGS.child_optim_algo == "adam":
          optimizer = optim.Adam(
              child_model.parameters(),
              eps=1e-3,
              weight_decay=FLAGS.child_l2_reg,
          )  # with L2
      else:
          raise ValueError(
              "Unknown optim_algo {}".format(FLAGS.child_optim_algo)
          )
  
      child_model.to(device)
      criterion.to(device)
  
      logger.info("Start training")
      start_time = time.time()
      step = 0
  
      # save path
      model_save_path = os.path.join(FLAGS.output_dir, "model.pth")
      best_model_save_path = os.path.join(
          FLAGS.output_dir, "best_model.pth"
      )
      best_acc = 0
      start_epoch = 0
      if FLAGS.load_checkpoint:
          if os.path.isfile(model_save_path):
              checkpoint = torch.load(
                  model_save_path, map_location=torch.device("cpu")
              )
              step = checkpoint["step"]
              start_epoch = checkpoint["epoch"]
              child_model.load_state_dict(
                  checkpoint["child_model_state_dict"]
              )
              optimizer.load_state_dict(
                  checkpoint["optimizer_state_dict"]
              )
  
      for epoch in range(start_epoch, FLAGS.num_epochs):
          lr = update_lr(
              optimizer,
              epoch,
              l2_reg=FLAGS.child_l2_reg,
              lr_warmup_val=None,
              lr_init=FLAGS.child_lr,
              lr_decay_scheme=FLAGS.child_lr_decay_scheme,
              lr_max=FLAGS.child_lr_max,
              lr_min=FLAGS.child_lr_min,
              lr_T_0=FLAGS.child_lr_T_0,
              lr_T_mul=FLAGS.child_lr_T_mul,
          )
          child_model.train()
          for batch in train_dataloader:
              (sent_ids, mask), labels = batch
  
              sent_ids = sent_ids.to(device, non_blocking=True)
              mask = mask.to(device, non_blocking=True)
              labels = labels.to(device, non_blocking=True)
  
              step += 1
  
              logits = child_model((sent_ids, mask))  # run
  
              loss = criterion(logits, labels.long())
              loss = loss.mean()
              preds = logits.argmax(dim=1).long()
              acc = torch.eq(preds, labels.long()).long().sum().item()
  
              optimizer.zero_grad()
              loss.backward()
              grad_norm = 0
              trainable_params = child_model.parameters()
  
              assert (
                  FLAGS.child_grad_bound is not None
              ), "Need grad_bound to clip gradients."
              # compute the gradient norm value
              grad_norm = nn.utils.clip_grad_norm_(
                  trainable_params, 99999999
              )
              for param in trainable_params:
                  nn.utils.clip_grad_norm_(
                      param, FLAGS.child_grad_bound
                  )  # clip grad
  
              optimizer.step()
  
              if step % FLAGS.log_every == 0:
                  curr_time = time.time()
                  log_string = ""
                  log_string += "epoch={:<6d}".format(epoch)
                  log_string += "ch_step={:<6d}".format(step)
                  log_string += " loss={:<8.6f}".format(loss)
                  log_string += " lr={:<8.4f}".format(lr)
                  log_string += " |g|={:<8.4f}".format(grad_norm)
                  log_string += " tr_acc={:<3d}/{:>3d}".format(
                      acc, logits.size()[0]
                  )
                  log_string += " mins={:<10.2f}".format(
                      float(curr_time - start_time) / 60
                  )
                  logger.info(log_string)
  
          epoch += 1
          save_state = {
              "step": step,
              "epoch": epoch,
              "child_model_state_dict": child_model.state_dict(),
              "optimizer_state_dict": optimizer.state_dict(),
          }
          torch.save(save_state, model_save_path)
          child_model.eval()
          logger.info("Epoch {}: Eval".format(epoch))
          eval_acc, eval_loss = eval_once(
              child_model,
              device,
              "test",
              criterion,
              test_dataloader=test_dataloader,
          )
          logger.info(
              "ch_step={} {}_accuracy={:<6.4f} {}_loss={:<6.4f}".format(
                  step, "test", eval_acc, "test", eval_loss
              )
          )
          if eval_acc > best_acc:
              best_acc = eval_acc
              logger.info("Save best model")
              save_state = {
                  "step": step,
                  "epoch": epoch,
                  "child_model_state_dict": child_model.state_dict(),
                  "optimizer_state_dict": optimizer.state_dict(),
              }
              torch.save(save_state, best_model_save_path)
  
      return eval_acc
  
imports: import sys
  import os
  import logging
  import pickle
  import shutil
  import random
  import math
  import time
  import datetime
  import argparse
  import distutils.util
  import numpy
  import torch
  from torch import nn
  from torch import optim
  from torch.utils.data import DataLoader
  import torch.nn.functional
  from model import Model
  from nni.nas.pytorch.fixed import apply_fixed_architecture
  from dataloader import read_data_sst
label_desc: Use image node_id 8 for calling a global function with example usage: train(device, data_path, output_dir, num_layers) and returns: eval_acc