class_docstr: Training loop for general machine learning models.
  Based on the training loop from the objax CIFAR10 example code.
class_docstr_tok: ["Training","loop","for","general","machine","learning","models",".","Based","on","the","training","loop","from","the","objax","CIFAR10","example","code","."]
code: def train_step(
      self, summary: Summary, data: dict, progress: np.ndarray
  ):
      kv = self.train_op(
          progress, data["image"].numpy(), data["label"].numpy()
      )
      for k, v in kv.items():
          if jn.isnan(v):
              raise ValueError("NaN, try reducing learning rate", k)
          if summary is not None:
              summary.scalar(k, float(v))
  
imports: import functools
  import os
  import shutil
  from typing import Callable
  import json
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  import objax
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo import convnet, wide_resnet
  from dataset import DataSet
label_desc: Use image node_id 2 for calling the TrainLoop obj's underlying member method code with example usage: obj.train_step(summary, data, progress) without return types