code: def run_customized_training(
      strategy,
      bert_config,
      max_seq_length,
      max_predictions_per_seq,
      model_dir,
      steps_per_epoch,
      steps_per_loop,
      epochs,
      initial_lr,
      input_files,
      train_batch_size,
  ):
      """Run BERT pretrain model training using low-level API."""
      if strategy is not None:
          num_replicas_in_sync = strategy.num_replicas_in_sync
      else:
          num_replicas_in_sync = 1
  
      train_input_fn = get_pretrain_dataset_fn(
          input_files,
          max_seq_length,
          max_predictions_per_seq,
          train_batch_size,
          num_replicas_in_sync,
      )
  
      def _get_pretrain_model():
          """Gets a pretraining model."""
          pretrain_model, core_model = bert_models.pretrain_model(
              bert_config, max_seq_length, max_predictions_per_seq
          )
  
          pretrain_model.optimizer = tf.optimizers.Adam(lr=initial_lr)
          if FLAGS.fp16_implementation == "graph_rewrite":
              pretrain_model.optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(
                  pretrain_model.optimizer
              )
          return pretrain_model, core_model
  
      time_callback = keras_utils.TimeHistory(
          train_batch_size * steps_per_loop, 1
      )
  
      ##########################################################################
      # Build with Graph mode and AutoDist scope in bert_utils
      trained_model = bert_utils.run_customized_training_loop(
          strategy=strategy,
          model_fn=_get_pretrain_model,
          loss_fn=get_loss_fn(
              loss_factor=1.0 / num_replicas_in_sync
              if FLAGS.scale_loss
              else 1.0
          ),
          model_dir=model_dir,
          train_input_fn=train_input_fn,
          steps_per_epoch=steps_per_epoch,
          steps_per_loop=steps_per_loop,
          epochs=epochs,
          sub_model_export_name="pretrained/bert_model",
          custom_callbacks=[time_callback],
      )
      ##########################################################################
      return trained_model
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import yaml
  import os
  import sys
  import tensorflow
  from absl import app
  from absl import flags
  from absl import logging
  from utils.logs import logger
  from utils.misc import keras_utils
  from utils import bert_modeling
  from utils import bert_models
  from utils import common_flags
  from utils import input_pipeline
  from utils import bert_utils
  from autodist import AutoDist
  from autodist.strategy.all_reduce_strategy import AllReduce
  from autodist.strategy.ps_strategy import PS
  from autodist.strategy.ps_lb_strategy import PSLoadBalancing
  from autodist.strategy.parallax_strategy import Parallax
  from autodist.strategy.partitioned_ps_strategy import PartitionedPS
label_desc: Use image node_id 3 for calling a global function with example usage: run_customized_training(strategy, bert_config, max_seq_length, max_predictions_per_seq, model_dir, steps_per_epoch, steps_per_loop, epochs, initial_lr, input_files, train_batch_size) and returns: trained_model, pretrain_model,  core_model