code: def main():
      parser = argparse.ArgumentParser(
          description="PyTorch Tensorboard Example"
      )
      parser.add_argument(
          "--cluster_mode",
          type=str,
          default="local",
          help="The cluster mode, such as local, yarn, spark-submit or k8s.",
      )
      parser.add_argument(
          "--runtime",
          type=str,
          default="spark",
          help="The runtime backend, one of spark or ray.",
      )
      parser.add_argument(
          "--address",
          type=str,
          default="",
          help="The cluster address if the driver connects to an existing "
          "ray cluster. If it is empty, a new Ray cluster will be created.",
      )
      parser.add_argument(
          "--backend",
          type=str,
          default="spark",
          help="The backend of PyTorch Estimator; "
          "spark, ray and bigdl are supported.",
      )
      parser.add_argument(
          "--batch_size",
          type=int,
          default=4,
          help="The training batch size",
      )
      parser.add_argument(
          "--epochs",
          type=int,
          default=2,
          help="The number of epochs to train for",
      )
      parser.add_argument(
          "--data_dir",
          type=str,
          default="./data",
          help="The path of dataset",
      )
      parser.add_argument(
          "--download",
          type=bool,
          default=True,
          help="Download dataset or not",
      )
      args = parser.parse_args()
  
      if args.runtime == "ray":
          init_orca_context(runtime=args.runtime, address=args.address)
      else:
          if args.cluster_mode == "local":
              init_orca_context()
          elif args.cluster_mode.startswith("yarn"):
              init_orca_context(
                  cluster_mode=args.cluster_mode, cores=4, num_nodes=2
              )
          elif args.cluster_mode == "spark-submit":
              init_orca_context(cluster_mode=args.cluster_mode)
  
      tensorboard_dir = args.data_dir + "runs"
      writer = SummaryWriter(
          tensorboard_dir + "/fashion_mnist_experiment_1"
      )
      # constant for classes
      classes = (
          "T-shirt/top",
          "Trouser",
          "Pullover",
          "Dress",
          "Coat",
          "Sandal",
          "Shirt",
          "Sneaker",
          "Bag",
          "Ankle Boot",
      )
  
      # plot some random training images
      dataiter = iter(
          train_data_creator(
              config={},
              batch_size=4,
              download=args.download,
              data_dir=args.data_dir,
          )
      )
      images, labels = dataiter.next()
  
      # create grid of images
      img_grid = torchvision.utils.make_grid(images)
  
      # show images
      matplotlib_imshow(img_grid, one_channel=True)
  
      # write to tensorboard
      writer.add_image("four_fashion_mnist_images", img_grid)
  
      # inspect the model using tensorboard
      writer.add_graph(model_creator(config={}), images)
      writer.close()
  
      # training loss vs. epochs
      criterion = nn.CrossEntropyLoss()
      batch_size = args.batch_size
      epochs = args.epochs
      if args.backend in ["ray", "spark"]:
          orca_estimator = Estimator.from_torch(
              model=model_creator,
              optimizer=optimizer_creator,
              loss=criterion,
              metrics=[Accuracy()],
              model_dir=os.getcwd(),
              use_tqdm=True,
              backend=args.backend,
          )
          stats = orca_estimator.fit(
              train_data_creator, epochs=epochs, batch_size=batch_size
          )
  
          for stat in stats:
              writer.add_scalar(
                  "training_loss", stat["train_loss"], stat["epoch"]
              )
          print("Train stats: {}".format(stats))
          val_stats = orca_estimator.evaluate(
              validation_data_creator, batch_size=batch_size
          )
          print("Validation stats: {}".format(val_stats))
          orca_estimator.shutdown()
      else:
          invalidInputError(
              False,
              "Only ray, and spark are supported "
              "as the backend, but got {}".format(args.backend),
          )
  
      stop_orca_context()
  
imports: from __future__ import print_function
  import os
  import argparse
  import numpy
  import matplotlib.pyplot
  import torch
  import torchvision
  import torchvision.transforms
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.utils.tensorboard import SummaryWriter
  from bigdl.dllib.utils.log4Error import invalidInputError
  from bigdl.orca import init_orca_context, stop_orca_context
  from bigdl.orca.learn.pytorch import Estimator
  from bigdl.orca.learn.metrics import Accuracy
  from bigdl.orca.learn.trigger import EveryEpoch
label_desc: Use image node_id 6 for calling a global function with example usage: main() without return types