code: def train_step(x_batch, y_batch):
      # forward + backward
      with tf.GradientTape() as tape:
          ## compute outputs
          _logits = vgg(x_batch)
          ## compute loss and update model
          _loss = loss_object(_logits, y_batch)
  
      grad = tape.gradient(_loss, train_weights)
      optimizer.apply_gradients(zip(grad, train_weights))
      return _loss
  
imports: import os
  import time
  import psutil
  import tensorflow
  import tensorlayer
  from exp_config import BATCH_SIZE, LERANING_RATE, MONITOR_INTERVAL, NUM_ITERS, random_input_generator
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(x_batch, y_batch) and returns: _loss