code: def train(
      self,
      num_train_epochs: int,
      train_size: int,
      train: DataSet,
      test: DataSet,
      logdir: str,
  ):
      checkpoint = objax.io.Checkpoint(
          logdir, keep_ckpts=5, makedir=True
      )
      start_epoch, last_ckpt = checkpoint.restore(self.vars())
      train_iter = iter(train)
      progress = np.zeros(
          jax.local_device_count(), "f"
      )  # for multi-GPU
  
      with SummaryWriter(os.path.join(logdir, "tb")) as tensorboard:
          for epoch in range(start_epoch, num_train_epochs):
              with self.vars().replicate():
                  # Train
                  summary = Summary()
                  loop = trange(
                      0,
                      train_size,
                      self.params.batch,
                      leave=False,
                      unit="img",
                      unit_scale=self.params.batch,
                      desc="Epoch %d/%d"
                      % (1 + epoch, num_train_epochs),
                  )
                  for step in loop:
                      progress[:] = (step + (epoch * train_size)) / (
                          num_train_epochs * train_size
                      )
                      self.train_step(
                          summary, next(train_iter), progress
                      )
  
                  # Eval
                  accuracy, total = 0, 0
                  for data in tqdm(
                      test, leave=False, desc="Evaluating"
                  ):
                      total += data["image"].shape[0]
                      preds = np.argmax(
                          self.predict(data["image"].numpy()), axis=1
                      )
                      accuracy += (preds == data["label"].numpy()).sum()
                  accuracy /= total
                  summary.scalar("eval/accuracy", 100 * accuracy)
                  print(
                      "Epoch %04d  Loss %.2f  Accuracy %.2f"
                      % (
                          epoch + 1,
                          summary["losses/xe"](),
                          summary["eval/accuracy"](),
                      )
                  )
                  tensorboard.write(
                      summary, step=(epoch + 1) * train_size
                  )
  
              checkpoint.save(self.vars(), epoch + 1)
  
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 3 for calling the TrainLoop obj's underlying member method code with example usage: obj.train(num_train_epochs, train_size, train, test, logdir) without return types