code: def _train_epoch(
      epoch,
      model: centernet_vid.CenterNetVID,
      dataset,
      logger,
      number_of_steps_previous_epochs,
      max_num_steps_epoch,
      input_image_size,
  ):
      """Train one epoch."""
      strategy = tf.distribute.get_strategy()
  
      if input_image_size:
          print(input_image_size)
  
      def distributed_step(sample):
          outputs = {}
          logging.info(sample["name"])
          per_replica_outputs = strategy.run(
              model.train_sample, args=(sample,)
          )
          losses_value = {}
          for key, value in per_replica_outputs.items():
              losses_value[key] = strategy.reduce(
                  tf.distribute.ReduceOp.SUM, value, axis=None
              )
  
          losses_val_value = {
              "total_loss": tf.constant(0, dtype=tf.float32)
          }
          total_loss_diff = tf.abs(
              losses_value["total_loss"]
              - losses_val_value["total_loss"]
          )
          return {
              **losses_value,
              "total_loss_diff": total_loss_diff,
          }, outputs
  
      if FLAGS.run_graph:
          distributed_step = tf.function(distributed_step)
  
      logger.reset_losses()
      dataset_iterator = iter(dataset)
  
      n_steps = 0
      while True:
          start_time = time.time()
          sample = tf_utils.get_next_sample_dataset(dataset_iterator)
  
          if (
              sample is None
              or max_num_steps_epoch is not None
              and n_steps > max_num_steps_epoch
          ):
              break
          n_steps += FLAGS.batch_size
  
          total_num_steps = n_steps + number_of_steps_previous_epochs
          logger.record_scalar(
              "meta/time_read",
              time.time() - start_time,
              total_num_steps,
          )
          logger.record_scalar(
              "meta/learning_rate",
              model.optimizer.learning_rate(total_num_steps),
              total_num_steps,
          )
          start_time = time.time()
          losses_value, _ = distributed_step(sample)
          logging.info(
              "Loss: %f/%f \t LR: %f GD: %f",
              losses_value["total_loss"].numpy(),
              losses_value["total_loss_diff"].numpy(),
              model.optimizer.learning_rate(total_num_steps),
              losses_value["gradients_norm"].numpy(),
          )
          logger.record_scalar(
              "meta/forward_pass",
              time.time() - start_time,
              total_num_steps,
          )
          logger.record_losses(
              "iterations/", losses_value, total_num_steps
          )
  
      logger.record_losses_epoch("epoch/", epoch)
      return n_steps
  
imports: import os
  import pickle
  import time
  from absl import app
  from absl import flags
  from absl import logging
  import matplotlib.pyplot
  import numpy
  import tensorflow
  import tensorflow_graphics.projects.points_to_3Dobjects.data_preparation.extract_protos
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import collision_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import cross_entropy_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import focal_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import focal_loss_sparse
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import learning_rate_schedule
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import regression_huber_loss
  from tensorflow_graphics.projects.points_to_3Dobjects.losses import regression_l1_loss
  import tensorflow_graphics.projects.points_to_3Dobjects.models.centernet_vid
  from tensorflow_graphics.projects.points_to_3Dobjects.transforms import transforms_factory
  from tensorflow_graphics.projects.points_to_3Dobjects.utils import evaluator
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.io
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.logger
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.plot
  import tensorflow_graphics.projects.points_to_3Dobjects.utils.tf_utils
  from google3.pyglib import gfile
label_desc: Use image node_id 7 for calling a global function with example usage: _train_epoch(epoch, model, dataset, logger, number_of_steps_previous_epochs, max_num_steps_epoch, input_image_size) and returns: n_steps, dict,  outputs