code: def main():
      parser = argparse.ArgumentParser(
          description="PyTorch MNIST Example"
      )
      parser.add_argument(
          "--batch-size",
          type=int,
          default=64,
          metavar="N",
          help="input batch size for training (default: 64)",
      )
      parser.add_argument(
          "--test-batch-size",
          type=int,
          default=1000,
          metavar="N",
          help="input batch size for testing (default: 1000)",
      )
      parser.add_argument(
          "--epochs",
          type=int,
          default=20,
          metavar="N",
          help="number of epochs to train",
      )
      parser.add_argument(
          "--lr",
          type=float,
          default=1.0,
          metavar="LR",
          help="learning rate",
      )
      parser.add_argument(
          "--gamma",
          type=float,
          default=0.7,
          metavar="M",
          help="Learning rate step gamma (default: 0.7)",
      )
      parser.add_argument(
          "--dry-run",
          action="store_true",
          default=False,
          help="quickly check a single pass",
      )
      parser.add_argument(
          "--log-interval",
          type=int,
          default=10,
          metavar="N",
          help="how many batches to wait before logging training status",
      )
      parser.add_argument(
          "--no-moe",
          action="store_true",
          default=False,
          help="if disabling moe layer and using ffn layer instead",
      )
      args = parser.parse_args()
  
      use_cuda = torch.cuda.is_available()
      device = penv.local_device
      torch.manual_seed(1)
  
      train_kwargs = {"batch_size": args.batch_size}
      test_kwargs = {"batch_size": args.test_batch_size}
      if use_cuda:
          cuda_kwargs = {
              "num_workers": 1,
              "pin_memory": True,
              "shuffle": True,
          }
          train_kwargs.update(cuda_kwargs)
          test_kwargs.update(cuda_kwargs)
  
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
  
      if int(torch.os.environ.get("LOCAL_RANK", 0)) == 0:
          dataset1 = Net.DATASET_TARGET(
              "/tmp/data",
              train=True,
              download=True,
              transform=transform,
          )
          dataset2 = Net.DATASET_TARGET(
              "/tmp/data", train=False, transform=transform
          )
          net.barrier()
      else:
          net.barrier()
          dataset1 = Net.DATASET_TARGET(
              "/tmp/data",
              train=True,
              download=False,
              transform=transform,
          )
          dataset2 = Net.DATASET_TARGET(
              "/tmp/data", train=False, transform=transform
          )
      net.barrier()
  
      train_loader = torch.utils.data.DataLoader(
          dataset1, **train_kwargs
      )
      test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
  
      model = Net(use_moe=not args.no_moe).to(device)
      optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
  
      penv.dist_print(
          "Model = %s.\nShared parameter items = %d (as replicas), expert parameter items = %d (as local x %d devices)."
          % (
              model,
              len(
                  [
                      x
                      for x in model.parameters()
                      if not hasattr(x, "skip_allreduce")
                  ]
              ),
              len(
                  [
                      x
                      for x in model.parameters()
                      if hasattr(x, "skip_allreduce")
                  ]
              ),
              penv.global_size,
          )
      )
  
      scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
      torch.manual_seed(penv.global_rank + 1)
  
      peak_accuracy = 0
      for epoch in range(1, args.epochs + 1):
          train(args, model, device, train_loader, optimizer, epoch)
          peak_accuracy = max(
              peak_accuracy, test(model, device, test_loader)
          )
          scheduler.step()
  
      penv.dist_print(
          "Peak validation accuracy = {:.2f}%".format(peak_accuracy)
      )
  
file_path: tutel/tutel/examples/moe_cifar10.py
imports: import argparse
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  from torch.optim.lr_scheduler import StepLR
  from tutel import system
  from tutel import moe
  from tutel import net
  import logging
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types