code: def __init__(self, model: Callable, nclass: int, **kwargs):
      super().__init__(nclass, **kwargs)
      self.model = model(3, nclass)
      self.opt = objax.optimizer.Momentum(self.model.vars())
      self.model_ema = objax.optimizer.ExponentialMovingAverageModule(
          self.model, momentum=0.999, debias=True
      )
  
      @objax.Function.with_vars(self.model.vars())
      def loss(x, label):
          logit = self.model(x, training=True)
          loss_wd = 0.5 * sum(
              (v.value**2).sum()
              for k, v in self.model.vars().items()
              if k.endswith(".w")
          )
          loss_xe = objax.functional.loss.cross_entropy_logits(
              logit, label
          ).mean()
          return loss_xe + loss_wd * self.params.weight_decay, {
              "losses/xe": loss_xe,
              "losses/wd": loss_wd,
          }
  
      gv = objax.GradValues(loss, self.model.vars())
  
      @objax.Function.with_vars(self.vars())
      def train_op(progress, x, y):
          g, v = gv(x, y)
          lr = self.params.lr * jn.cos(progress * (7 * jn.pi) / (2 * 8))
          self.opt(lr, objax.functional.parallel.pmean(g))
          self.model_ema.update_ema()
          return objax.functional.parallel.pmean(
              {"monitors/lr": lr, **v[1]}
          )
  
      self.predict = objax.Parallel(
          objax.nn.Sequential(
              [
                  objax.ForceArgs(self.model_ema, training=False),
                  objax.functional.softmax,
              ]
          )
      )
      self.train_op = objax.Parallel(train_op, reduce=lambda x: x[0])
  
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 TrainModule object from inherited base classes: TrainLoop with example: obj = TrainModule(model, nclass)