code: def compute_gradients(model, images, labels, num_replicas=1):
      with tf.GradientTape() as grad_tape:
          logits = model(images, training=True)
          loss = tf.compat.v1.losses.softmax_cross_entropy(
              logits=logits, onehot_labels=labels
          )
          tf.compat.v2.summary.write("loss", loss)
          if num_replicas != 1:
              loss /= num_replicas
  
      # TODO(b/110991947): We can mistakenly trace the gradient call in
      # multi-threaded environment. Explicitly disable recording until
      # this is fixed.
      with record.stop_recording():
          grads = grad_tape.gradient(loss, model.variables)
      return grads
  
imports: import gc
  import os
  import tempfile
  import time
  import tensorflow
  from tensorflow.python.client import device_lib
  from tensorflow.python.eager import context
  from tensorflow.python.eager import record
  from tensorflow.python.eager.benchmarks.resnet50 import resnet50
  from tensorflow.python.eager.benchmarks.resnet50 import resnet50_test_util
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: compute_gradients(model, images, labels, num_replicas) and returns: grads