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=14,
          metavar="N",
          help="number of epochs to train (default: 14)",
      )
      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",
      )
      parser.add_argument(
          "--save-model",
          action="store_true",
          default=False,
          help="For Saving the current model",
      )
  
      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,)),
          ]
      )
      dataset1 = datasets.MNIST(
          "../data", train=True, download=True, transform=transform
      )
      dataset2 = datasets.MNIST(
          "../data", train=False, transform=transform
      )
      train_loader = torch.utils.data.DataLoader(
          dataset1, **train_kwargs
      )
      test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
  
      model = Net().to(device)
      scripted_model = torch.jit.script(model)  # scripting the model
      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,
              scripted_model,
              device,
              train_loader,
              optimizer,
              epoch,
          )
          scheduler.step()
      test(scripted_model, device, test_loader)
  
      # Saving model and running inference with BentoML:
  
      # Option1: save natively with bentoml.torchscript_iris
      bentoml.torchscript.save_model(
          "torchscript_mnist",
          scripted_model,
          signatures={"__call__": {"batchable": True}},
      )
      model_runner = bentoml.torchscript.get(
          "torchscript_mnist"
      ).to_runner()
      model_runner.init_local()
  
      test_datapoint, test_target = next(iter(test_loader))
      prediction = model_runner.run(
          test_datapoint[0].unsqueeze(0).numpy()
      )
      actual = test_target[0].item()
      predicted = np.argmax(prediction).item()
      print(
          "\nPREDICTION RESULT: ACTUAL: {}, PREDICTED: {}".format(
              str(actual), str(predicted)
          )
      )
  
      # Option2: save MLflow model and import MLflow pyfunc model to BentoML
      with mlflow.start_run():
          # logging scripted model
          mlflow.pytorch.log_model(scripted_model, "model")
  
          # Import logged mlflow model to BentoML model store for serving:
          model_uri = mlflow.get_artifact_uri("model")
          bento_model = bentoml.mlflow.import_model(
              "mlflow_torch_mnist",
              model_uri,
              signatures={"predict": {"batchable": True}},
          )
          print(f"Model imported to BentoML: {bento_model}")
  
          model_runner = bentoml.mlflow.get(
              "mlflow_torch_mnist"
          ).to_runner()
          model_runner.init_local()
  
          test_datapoint, test_target = next(iter(test_loader))
          prediction = model_runner.predict.run(
              test_datapoint[0].unsqueeze(0).numpy()
          )
          actual = test_target[0].item()
          predicted = np.argmax(prediction).item()
          print(
              "\nPREDICTION RESULT: ACTUAL: {}, PREDICTED: {}".format(
                  str(actual), str(predicted)
              )
          )
  
file_path: BentoML/examples/mlflow/torchscript/MNIST/mnist_torchscript.py
imports: import argparse
  import mlflow
  import mlflow.pytorch
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.optim.lr_scheduler import StepLR
  from torchvision import datasets
  from torchvision import transforms
  import bentoml
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types