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,
      )
  
      model = Net().to(device)
  
      save_checkpoint_dir = args["save_checkpoint_dir"]
      save_checkpoint_path = os.path.join(
          save_checkpoint_dir, "model.pth"
      )
      load_checkpoint_path = os.path.join(
          args["load_checkpoint_dir"], "model.pth"
      )
  
      if os.path.isfile(load_checkpoint_path):
          model_state_dict = load_checkpoint(load_checkpoint_path)
          logger.info("test : ", load_checkpoint_path)
          logger.info(type(model_state_dict))
          model.load_state_dict(model_state_dict)
  
      optimizer = optim.SGD(
          model.parameters(), lr=args["lr"], momentum=args["momentum"]
      )
  
      # epoch is perturbation interval
      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.")
  
      if not os.path.exists(save_checkpoint_dir):
          os.makedirs(save_checkpoint_dir)
      save_checkpoint(model, save_checkpoint_path)
  
file_path: nni/examples/trials/mnist-pbt-tuner-pytorch/mnist.py
imports: import argparse
  import logging
  import os
  import nni
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types