code: def test_train_epoch(self):
      model = snt.Sequential(
          [
              snt.Flatten(),
              snt.Linear(10),
          ]
      )
  
      optimizer = snt.optimizers.SGD(0.1)
  
      dataset = (
          tf.data.Dataset.from_tensor_slices(
              (
                  tf.random.normal([2, 8, 8, 1]),
                  tf.ones([2], dtype=tf.int64),
              )
          )
          .batch(2)
          .repeat(4)
      )
  
      for _ in range(3):
          loss = simple_mnist.train_epoch(model, optimizer, dataset)
      self.assertEqual(loss.shape, [])
      self.assertEqual(loss.dtype, tf.float32)
  
file_path: sonnet/examples/simple_mnist_test.py
imports: import sonnet
  from examples import simple_mnist
  from sonnet.src import test_utils
  import tensorflow
label_desc: Use image node_id 2 for calling the SimpleMnistTest obj's underlying member method code with example usage: obj.test_train_epoch() without return types