code: def training_step(images, labels, first_batch):
      with tf.GradientTape() as tape:
          probs = mnist_model(images, training=True)
          loss_value = loss(labels, probs)
  
      tape = bps.DistributedGradientTape(tape)
  
      grads = tape.gradient(loss_value, mnist_model.trainable_variables)
      opt.apply_gradients(zip(grads, mnist_model.trainable_variables))
  
      # Note: broadcast should be done after the first gradient step to ensure optimizer
      # initialization.
      if first_batch:
          bps.broadcast_variables(mnist_model.variables, root_rank=0)
          bps.broadcast_variables(opt.variables(), root_rank=0)
  
      return loss_value
  
label_desc: Use image node_id 1 for calling a global function with example usage: training_step(images, labels, first_batch) and returns: loss_value