code: def main():
      # Training settings
      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=5,
          metavar="N",
          help="number of epochs to train (default: 5)",
      )
      parser.add_argument(
          "--lr",
          type=float,
          default=1.0,
          metavar="LR",
          help="learning rate (default: 1.0)",
      )
      parser.add_argument(
          "--gamma",
          type=float,
          default=0.7,
          metavar="M",
          help="Learning rate step gamma (default: 0.7)",
      )
      parser.add_argument(
          "--no-cuda",
          action="store_true",
          default=False,
          help="disables CUDA training",
      )
      parser.add_argument(
          "--dry-run",
          action="store_true",
          default=False,
          help="quickly check a single pass",
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=1,
          metavar="S",
          help="random seed (default: 1)",
      )
      parser.add_argument(
          "--log-interval",
          type=int,
          default=10,
          metavar="N",
          help="how many batches to wait before logging training status",
      )
      args = parser.parse_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")
  
      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,)),
          ]
      )
      train_ds = datasets.MNIST(
          "data", train=True, download=True, transform=transform
      )
      test_ds = datasets.MNIST("data", train=False, transform=transform)
      train_loader = DataLoader(train_ds, **train_kwargs)
      test_loader = DataLoader(test_ds, **test_kwargs)
  
      model = net.CNN().to(device)
      optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
  
      scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
      for epoch in range(1, args.epochs + 1):
          train(args, model, device, train_loader, optimizer, epoch)
          test(model, device, test_loader)
          scheduler.step()
  
      import bentoml
  
      model = bentoml.pytorch.save_model(
          "mnist_cnn",
          model,
          signatures={"__call__": {"batchable": True}},
          external_modules=[net],
      )
      print(f"Saved: {model}")
  
file_path: BentoML/examples/custom_model_runner/train.py
imports: from __future__ import print_function
  import argparse
  import net
  import torch
  import torch.nn.functional
  import torch.optim
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data import DataLoader
  from torchvision import datasets
  from torchvision import transforms
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types