code: def test_tf_nano_multiprocessing_customized_loop():
      from bigdl.nano.tf.keras import nano_multiprocessing, nano
  
      global_batch_size = 32
      model = tf.keras.Sequential(
          [tf.keras.layers.Dense(1, input_shape=(1,))]
      )
      optimizer = tf.keras.optimizers.SGD()
  
      dataset = (
          tf.data.Dataset.from_tensors(([1.0], [1.0]))
          .repeat(128)
          .batch(global_batch_size)
      )
  
      loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)
  
      @nano_multiprocessing
      @tf.function
      def train_step(inputs, model, loss_object, optimizer):
          features, labels = inputs
  
          with tf.GradientTape() as tape:
              predictions = model(features, training=True)
              loss = loss_object(labels, predictions)
  
          gradients = tape.gradient(loss, model.trainable_variables)
          optimizer.apply_gradients(
              zip(gradients, model.trainable_variables)
          )
          return loss
  
      @nano(num_processes=2)
      def train_whole_data(
          model, dataset, loss_object, optimizer, train_step
      ):
          for inputs in dataset:
              print(train_step(inputs, model, loss_object, optimizer))
  
      train_whole_data(
          model, dataset, loss_object, optimizer, train_step
      )
  
file_path: BigDL/python/nano/test/tf/train/test_tf_nano_decorator.py
imports: import tensorflow
  from tensorflow import keras
  from tensorflow.keras import layers
  from tensorflow.keras import backend
  from bigdl.nano.tf.keras import Sequential, nano_bf16
label_desc: Use image node_id 2 for calling a global function with example usage: test_tf_nano_multiprocessing_customized_loop() and returns: loss