code: def run(
      train_batch_size,
      val_batch_size,
      epochs,
      lr,
      momentum,
      log_interval,
      log_dir,
  ):
      train_loader, val_loader = get_data_loaders(
          train_batch_size, val_batch_size
      )
      model = Net()
      writer = SummaryWriter(log_dir=log_dir)
  
      # Use TPU device
      device = xm.xla_device()
  
      model.to(device)  # Move model before creating optimizer
      optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
      criterion = nn.NLLLoss()
  
      # Create trainer and evaluator
      trainer = create_supervised_trainer(
          model,
          optimizer,
          criterion,
          device=device,
          output_transform=lambda x, y, y_pred, loss: [loss.item()],
      )
  
      val_metrics = {"accuracy": Accuracy(), "nll": Loss(criterion)}
      evaluator = create_supervised_evaluator(
          model, metrics=val_metrics, device=device
      )
  
      tracker = xm.RateTracker()
  
      # Add RateTracker as an output of the training step
      @trainer.on(Events.ITERATION_COMPLETED)
      def add_rate_tracker(engine):
          tracker.add(len(engine.state.batch))
          engine.state.output.append(tracker.global_rate())
  
      # Setup output values of the training step as EMA metrics
      RunningAverage(output_transform=lambda x: x[0]).attach(
          trainer, "batch_loss"
      )
      RunningAverage(output_transform=lambda x: x[1]).attach(
          trainer, "global_rate"
      )
  
      # Let's log the EMA metrics every `log_interval` iterations
      @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
      def log_training_loss(engine):
          writer.add_scalar(
              "training/batch_loss",
              engine.state.metrics["batch_loss"],
              engine.state.iteration,
          )
          writer.add_scalar(
              "training/global_rate",
              engine.state.metrics["global_rate"],
              engine.state.iteration,
          )
  
      @trainer.on(Events.EPOCH_COMPLETED)
      def log_training_results(engine):
          evaluator.run(train_loader)
          metrics = evaluator.state.metrics
          avg_accuracy = metrics["accuracy"]
          avg_nll = metrics["nll"]
          print(
              f"Training Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
          )
          writer.add_scalar(
              "training/avg_loss", avg_nll, engine.state.epoch
          )
          writer.add_scalar(
              "training/avg_accuracy", avg_accuracy, engine.state.epoch
          )
  
      @trainer.on(Events.EPOCH_COMPLETED)
      def log_validation_results(engine):
          evaluator.run(val_loader)
          metrics = evaluator.state.metrics
          avg_accuracy = metrics["accuracy"]
          avg_nll = metrics["nll"]
          print(
              f"Validation Results - Epoch: {engine.state.epoch} Avg accuracy: {avg_accuracy:.2f} Avg loss: {avg_nll:.2f}"
          )
          writer.add_scalar(
              "valdation/avg_loss", avg_nll, engine.state.epoch
          )
          writer.add_scalar(
              "valdation/avg_accuracy", avg_accuracy, engine.state.epoch
          )
  
      # kick everything off
      trainer.run(train_loader, max_epochs=epochs)
  
      writer.close()
  
file_path: ignite/examples/mnist/mnist_with_tensorboard_on_tpu.py
imports: from argparse import ArgumentParser
  import torch.nn.functional
  from torch import nn
  from torch.optim import SGD
  from torch.utils.data import DataLoader
  from torch.utils.tensorboard import SummaryWriter
  from torchvision.datasets import MNIST
  from torchvision.transforms import Compose, Normalize, ToTensor
  from ignite.engine import create_supervised_evaluator, create_supervised_trainer, Events
  from ignite.metrics import Accuracy, Loss, RunningAverage
label_desc: Use image node_id 2 for calling a global function with example usage: run(train_batch_size, val_batch_size, epochs, lr, momentum, log_interval, log_dir) without return types