code: def __init__(self, nclass: int, **kwargs):
      self.nclass = nclass
      self.params = EasyDict(kwargs)
  
imports: import functools
  import os
  from typing import Callable
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  from tqdm import tqdm, trange
  import objax
  from examples.image_classification.tfdata.data import DataSet
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo import convnet, wide_resnet
label_desc: Use image node_id 1 to create a new TrainLoop object from inherited base classes: objax with example: obj = TrainLoop(nclass)