code: def train_step(model, optimizer, X, y):
      """Training loop implementation; the training loop is custom so that we can
      compute loss and accuracy for each batch as we train, inside the loop.
  
      """
      with tf.GradientTape() as tape:
          logits = model(X, training=True)
          loss = compute_loss(y, logits)
  
      accuracy = compute_accuracy(y, logits)
      grads = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(zip(grads, model.trainable_variables))
      return {"loss": loss, "accuracy": accuracy}
  
imports: from __future__ import division, print_function
  import getpass
  from typing import Dict, Optional
  import fs
  import tensorflow
  import tensorflow_datasets
  import tqdm
  import uv
  import uv.tensorflow.reporter
  import uv.util
  from absl import app
  from uv.fs.reporter import FSReporter
  from uv.mlflow.reporter import MLFlowReporter
  from trainer import cli
label_desc: Use image node_id 8 for calling a global function with example usage: train_step(model, optimizer, X, y) and returns: dict