code: def train_epoch(
      model: snt.Module,
      optimizer: snt.Optimizer,
      dataset: tf.data.Dataset,
  ) -> tf.Tensor:
      loss = 0.0
      for images, labels in dataset:
          loss = train_step(model, optimizer, images, labels)
      return loss
  
file_path: sonnet/examples/simple_mnist.py
imports: from typing import Dict
  from absl import app
  import sonnet
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: train_epoch(model, optimizer, dataset) and returns: loss