code: def run_fn(fn_args: TrainerFnArgs):
      """Train the model based on given args.
  
      Args:
        fn_args: Holds args used to train the model as name/value pairs.
      """
      # Number of nodes in the first layer of the DNN
      first_dnn_layer_size = 100
      num_dnn_layers = 4
      dnn_decay_factor = 0.7
  
      tf_transform_output = tft.TFTransformOutput(
          fn_args.transform_output
      )
  
      train_dataset = _input_fn(
          fn_args.train_files,
          fn_args.data_accessor,
          tf_transform_output,
          40,
      )
      eval_dataset = _input_fn(
          fn_args.eval_files,
          fn_args.data_accessor,
          tf_transform_output,
          40,
      )
  
      # If no GPUs are found, CPU is used.
      mirrored_strategy = tf.distribute.MirroredStrategy()
      with mirrored_strategy.scope():
          model = _build_keras_model(
              # Construct layers sizes with exponetial decay
              hidden_units=[
                  max(
                      2,
                      int(first_dnn_layer_size * dnn_decay_factor**i),
                  )
                  for i in range(num_dnn_layers)
              ]
          )
  
      # Creating a custom callback using LambdaCallback to add batch data
      def batch_output(batch, logs):
          tf.summary.scalar("batch_loss", data=logs["loss"], step=batch)
          tf.summary.scalar(
              "batch_accuracy", data=logs["binary_accuracy"], step=batch
          )
          return batch
  
      batch_log_callback = LambdaCallback(on_batch_end=batch_output)
  
      # Write logs to path
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=fn_args.model_run_dir, update_freq="batch"
      )
  
      model.fit(
          train_dataset,
          steps_per_epoch=fn_args.train_steps,
          validation_data=eval_dataset,
          validation_steps=fn_args.eval_steps,
          callbacks=[tensorboard_callback, batch_log_callback],
      )
  
      # Export the model.
      export_serving_model(
          tf_transform_output, model, fn_args.serving_model_dir
      )
  
imports: from typing import List
  import absl
  from keras.callbacks import LambdaCallback
  import tensorflow
  import tensorflow_transform
  from tfx.components.trainer.executor import TrainerFnArgs
  from tfx.components.trainer.fn_args_utils import DataAccessor
  from tfx_bsl.tfxio import dataset_options
label_desc: Use image node_id 13 for calling a global function with example usage: run_fn(fn_args) and returns: batch