code: def __init__(self, nclass: int, model: Callable, **kwargs):
      super().__init__(nclass, **kwargs)
      self.model = model(nin=3, nclass=nclass, **kwargs)
      model_vars = self.model.vars()
      self.model_ema = objax.optimizer.ExponentialMovingAverageModule(
          self.model, momentum=0.999, debias=False
      )
      self.opt = objax.optimizer.Momentum(
          model_vars, momentum=0.9, nesterov=True
      )
  
      @objax.Function.with_vars(model_vars)
      def loss_function(x, u, y):
          xu = jn.concatenate(
              [x, u.reshape((-1,) + u.shape[2:])], axis=0
          )
          logit = self.model(xu, training=True)
          logit_x = logit[: x.shape[0]]
          logit_weak = logit[x.shape[0] :: 2]
          logit_strong = logit[x.shape[0] + 1 :: 2]
  
          xe = objax.functional.loss.cross_entropy_logits(
              logit_x, y
          ).mean()
  
          pseudo_labels = objax.functional.stop_gradient(
              objax.functional.softmax(logit_weak)
          )
          pseudo_mask = (
              pseudo_labels.max(axis=1) >= self.params.confidence
          ).astype(logit_weak.dtype)
          xeu = objax.functional.loss.cross_entropy_logits_sparse(
              logit_strong, pseudo_labels.argmax(axis=1)
          )
          xeu = (xeu * pseudo_mask).mean()
  
          wd = 0.5 * sum(
              objax.functional.loss.l2(v.value)
              for k, v in model_vars.items()
              if k.endswith(".w")
          )
  
          loss = xe + self.params.wu * xeu + self.params.wd * wd
          return loss, {
              "losses/xe": xe,
              "losses/xeu": xeu,
              "losses/wd": wd,
              "monitors/mask": pseudo_mask.mean(),
          }
  
      gv = objax.GradValues(loss_function, model_vars)
  
      @objax.Function.with_vars(self.vars())
      def train_op(step, x, u, y):
          g, v = gv(x, u, y)
          fstep = step[0] / (FLAGS.train_kimg << 10)
          lr = self.params.lr * jn.cos(fstep * (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.eval_op = objax.Parallel(
          objax.ForceArgs(self.model_ema, training=False)
      )
      self.train_op = objax.Parallel(train_op, reduce=lambda x: x[0])
  
file_path: objax/examples/fixmatch/fixmatch.py
imports: import os
  from typing import Callable
  import jax.numpy
  from absl import app, flags
  import objax
  from examples.fixmatch.libml.augment.augment import get_augment
  from examples.fixmatch.libml.data.fsl import DATASETS_LABELED
  from examples.fixmatch.libml.data.ssl import DATASETS_UNLABELED
  from examples.fixmatch.libml.models import ARCHS, network
  from examples.fixmatch.libml.train import TrainLoopSSL
  from examples.fixmatch.libml.util import setup_tf
label_desc: Use image node_id 1 to create a new FixMatch object from inherited base classes: TrainLoopSSL with example: obj = FixMatch(nclass, model)