code: def main(args):
      use_cuda = not args["no_cuda"] and torch.cuda.is_available()
  
      torch.manual_seed(args["seed"])
  
      device = torch.device("cuda" if use_cuda else "cpu")
  
      kwargs = (
          {"num_workers": 1, "pin_memory": True} if use_cuda else {}
      )
  
      data_dir = args["data_dir"]
  
      train_loader = torch.utils.data.DataLoader(
          datasets.MNIST(
              data_dir,
              train=True,
              download=True,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,)),
                  ]
              ),
          ),
          batch_size=args["batch_size"],
          shuffle=True,
          **kwargs,
      )
      test_loader = torch.utils.data.DataLoader(
          datasets.MNIST(
              data_dir,
              train=False,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,)),
                  ]
              ),
          ),
          batch_size=1000,
          shuffle=True,
          **kwargs,
      )
  
      hidden_size = args["hidden_size"]
  
      model = Net(hidden_size=hidden_size).to(device)
      get_and_apply_next_architecture(model)
      optimizer = optim.SGD(
          model.parameters(), lr=args["lr"], momentum=args["momentum"]
      )
  
      for epoch in range(1, args["epochs"] + 1):
          train(args, model, device, train_loader, optimizer, epoch)
          test_acc = test(args, model, device, test_loader)
  
          if epoch < args["epochs"]:
              # report intermediate result
              nni.report_intermediate_result(test_acc)
              logger.debug("test accuracy %g", test_acc)
              logger.debug("Pipe send intermediate result done.")
          else:
              # report final result
              nni.report_final_result(test_acc)
              logger.debug("Final result is %g", test_acc)
              logger.debug("Send final result done.")
  
imports: import os
  import argparse
  import logging
  from collections import OrderedDict
  import nni
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  from nni.nas.pytorch.mutables import LayerChoice, InputChoice
  from nni.algorithms.nas.pytorch.classic_nas import get_and_apply_next_architecture
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types