code: def _train_func(target_path, *args):
      mirrored_strategy = tf.distribute.MultiWorkerMirroredStrategy()
  
      actrual_args = [None] * len(args)
      new_model = None
  
      with mirrored_strategy.scope():
          for i, arg in enumerate(args):
              if arg[0] == "model":
                  actrual_args[i] = tf.keras.models.load_model(arg[1])
                  new_model = actrual_args[i]
                  continue
              # deserialize optimizer
              if arg[0] == "optimizer":
                  with open(arg[1], "rb") as f:
                      actrual_args[i] = cloudpickle.load(f)
                  continue
              # deserialize dataset
              if arg[0] == "dataset":
                  try:
                      from tensorflow.python.distribute.coordinator.values import (
                          deserialize_dataset_from_graph,
                      )
  
                      original_dataset = deserialize_dataset_from_graph(
                          arg[1], arg[2]
                      )
                  except ValueError:
                      original_dataset = tf.data.Dataset.from_generator(
                          arg[1], output_signature=arg[2]
                      )
                  actrual_args[
                      i
                  ] = mirrored_strategy.experimental_distribute_dataset(
                      original_dataset
                  )
                  continue
  
              if arg[0] == "loss":
                  with open(arg[1], "rb") as f:
                      original_loss_object = cloudpickle.load(f)
                      original_loss_object.reduction = (
                          tf.keras.losses.Reduction.NONE
                      )
  
                  def loss_object(*args, **kwargs):
                      per_example_loss = original_loss_object(
                          *args, **kwargs
                      )
                      if (
                          per_example_loss.shape == []
                          or per_example_loss.shape[0] == 0
                      ):
                          size = mirrored_strategy.num_replicas_in_sync
                          return (
                              tf.math.reduce_sum(per_example_loss)
                              / size
                          )
                      else:
                          size = (
                              per_example_loss.shape[0]
                              * mirrored_strategy.num_replicas_in_sync
                          )
                          return tf.nn.compute_average_loss(
                              per_example_loss, global_batch_size=size
                          )
  
                  actrual_args[i] = loss_object
                  continue
              # deserialize others
              if arg[0] == "others":
                  with open(arg[1], "rb") as f:
                      actrual_args[i] = cloudpickle.load(f)
                      if callable(actrual_args[i]) and isinstance(
                          actrual_args[i], nano_multiprocessing
                      ):
                          actrual_args[i] = partial(
                              actrual_args[i],
                              mirrored_strategy=mirrored_strategy,
                          )
  
          with open(target_path, "rb") as f:
              target_func = cloudpickle.load(f)
  
          try:
              res = target_func(*actrual_args)
          except TypeError:
              invalidInputError(
                  False,
                  "Please check if you have added 'nano_multiprocessing' "
                  "decorator to the train_step function.",
              )
  
          task_id = mirrored_strategy.cluster_resolver.task_id
          if task_id == 0:
              path = os.path.join("trained_model_weights")
              new_model.save_weights(path, overwrite=True)
  
      try:
          # Related to https://github.com/tensorflow/tensorflow/issues/50487, to avoid such error
          import atexit
  
          atexit.register(
              mirrored_strategy._extended._cross_device_ops._pool.close
          )
          atexit.register(
              mirrored_strategy._extended._host_cross_device_ops._pool.close
          )
      except AttributeError:
          pass
      return res
  
imports: import cloudpickle
  import tensorflow
  import numpy
  from functools import wraps, partial
  from tempfile import TemporaryDirectory
  import os
  import json
  from bigdl.nano.utils.common import schedule_processors
  from bigdl.nano.tf.keras.distributed_utils import _find_free_port
  from tensorflow.keras.losses import Loss
  import warnings
  from bigdl.nano.utils.common import invalidInputError
label_desc: Use image node_id 3 for calling a global function with example usage: _train_func(target_path) and returns: res, unknown, tf