code: def _multi_trial_test(epochs, batch_size, port, benchmark):
      # initalize dataset. Note that 50k+10k is used. It's a little different from paper
      transf = [
          transforms.RandomCrop(32, padding=4),
          transforms.RandomHorizontalFlip(),
      ]
      normalize = [
          transforms.ToTensor(),
          transforms.Normalize(
              [0.49139968, 0.48215827, 0.44653124],
              [0.24703233, 0.24348505, 0.26158768],
          ),
      ]
      train_dataset = serialize(
          CIFAR10,
          "data",
          train=True,
          download=True,
          transform=transforms.Compose(transf + normalize),
      )
      test_dataset = serialize(
          CIFAR10,
          "data",
          train=False,
          transform=transforms.Compose(normalize),
      )
  
      # specify training hyper-parameters
      training_module = NasBench101TrainingModule(max_epochs=epochs)
      # FIXME: need to fix a bug in serializer for this to work
      # lr_monitor = serialize(LearningRateMonitor, logging_interval='step')
      trainer = pl.Trainer(max_epochs=epochs, gpus=1)
      lightning = pl.Lightning(
          lightning_module=training_module,
          trainer=trainer,
          train_dataloader=pl.DataLoader(
              train_dataset, batch_size=batch_size, shuffle=True
          ),
          val_dataloaders=pl.DataLoader(
              test_dataset, batch_size=batch_size
          ),
      )
  
      strategy = Random()
  
      model = NasBench101()
  
      exp = RetiariiExperiment(model, lightning, [], strategy)
  
      exp_config = RetiariiExeConfig("local")
      exp_config.trial_concurrency = 2
      exp_config.max_trial_number = 20
      exp_config.trial_gpu_number = 1
      exp_config.training_service.use_active_gpu = False
  
      if benchmark:
          exp_config.benchmark = "nasbench101"
          exp_config.execution_engine = "benchmark"
  
      exp.run(exp_config, port)
  
imports: import click
  import nni
  import nni.retiarii.evaluator.pytorch.lightning
  import torch.nn
  import torchmetrics
  from nni.retiarii import model_wrapper, serialize
  from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
  from nni.retiarii.nn.pytorch import NasBench101Cell
  from nni.retiarii.strategy import Random
  from pytorch_lightning.callbacks import LearningRateMonitor
  from timm.optim import RMSpropTF
  from torch.optim.lr_scheduler import CosineAnnealingLR
  from torchvision import transforms
  from torchvision.datasets import CIFAR10
  from base_ops import Conv3x3BnRelu, Conv1x1BnRelu, Projection
label_desc: Use image node_id 1 for calling a global function with example usage: _multi_trial_test(epochs, batch_size, port, benchmark) without return types