code: def run_training_job():
      mnist_model = MnistModel()
  
      # Initialize a trainer
      logger = pl_loggers.TensorBoardLogger(args.log_path)
      trainer = Trainer(
          logger=logger,
          log_every_n_steps=1,
          max_epochs=args.epochs,
          enable_progress_bar=False,
          deterministic=True,  # Do we want a bit of noise?
          default_root_dir=args.log_path,
      )
      logger.save()
  
      print(f"Logging to path: {args.log_path}.")
  
      # Train the model and log time
      start = time.time()
      trainer.fit(model=mnist_model)
      end = time.time()
      train_time = end - start
  
      # Compute the validation accuracy
      val_accuracy = trainer.validate()[0]["val_acc"]
  
      # Log the number of model parameters
      num_params = trainer.model.num_params
  
      # Print outputs
      print(
          f"train time: {train_time}, val acc: {val_accuracy}, num_params: {num_params}"
      )
  
imports: import argparse
  import logging
  import os
  import time
  import warnings
  import torch
  from pytorch_lightning import LightningModule, loggers, Trainer
  from torch import nn
  from torch.nn import functional
  from torch.utils.data import DataLoader
  from torchmetrics.functional.classification.accuracy import multiclass_accuracy
  from torchvision import transforms
  from torchvision.datasets import MNIST
label_desc: Use image node_id 2 for calling a global function with example usage: run_training_job() without return types