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 __init__(self, nclass: int, **kwargs):
      self.nclass = nclass
      self.params = EasyDict(kwargs)
  
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 1 to create a new TrainLoop object from inherited base classes: objax with example: obj = TrainLoop(nclass)