code: def testDefunBackProp(self):
      convnet = cifar10_convnet.Cifar10ConvNet()
  
      @tf.function
      def do_training_step(image, labels):
          with tf.GradientTape() as tape:
              logits = convnet(image, is_training=True)["logits"]
              loss = tf.reduce_mean(
                  tf.nn.sparse_softmax_cross_entropy_with_logits(
                      logits=logits, labels=labels
                  )
              )
          grads = tape.gradient(loss, convnet.trainable_variables)
          return loss, grads
  
      image = tf.random.normal([4, 32, 32, 3])
      labels = np.random.randint(
          low=0, high=10, size=[4], dtype=np.int64
      )
      loss, grads = do_training_step(image, labels)
      self.assertEqual(loss.numpy().shape, ())
      for grad, var in zip(grads, convnet.trainable_variables):
          self.assertIsNotNone(grad)
          self.assertEqual(grad.numpy().shape, var.shape)
  
file_path: sonnet/sonnet/src/nets/cifar10_convnet_test.py
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import test_utils
  from sonnet.src.nets import cifar10_convnet
  import tensorflow
label_desc: Use image node_id 6 for calling the ModelTest obj's underlying member method code with example usage: obj.testDefunBackProp() and returns: loss,  grads