code: def main(argv):
      del argv
      # In this example we use tensorflow_datasets for loading cifar10, but you can use any dataset library you like.
      tf.config.experimental.set_visible_devices([], "GPU")
      DATA_DIR = os.path.join(os.environ["HOME"], "TFDS")
      data, info = tfds.load(
          name="cifar10",
          split="train",
          data_dir=DATA_DIR,
          with_info=True,
      )
      train_size = info.splits["train"].num_examples
      image_shape = info.features["image"].shape
      nclass = info.features["label"].num_classes
      train = DataSet.from_tfds(
          data, image_shape, augment_fn=lambda x: augment(x, 4)
      )
      test = DataSet.from_tfds(
          tfds.load(name="cifar10", split="test", data_dir=DATA_DIR),
          image_shape,
      )
      train = (
          train.cache()
          .shuffle(8192)
          .repeat()
          .parse()
          .augment()
          .batch(FLAGS.batch)
      )
      train = train.nchw().one_hot(nclass).prefetch(16)
      test = test.cache().parse().batch(FLAGS.batch).nchw().prefetch(16)
      del data, info
  
      # Define the network and train_it
      tm = TrainModule(
          network(FLAGS.arch),
          nclass=nclass,
          arch=FLAGS.arch,
          lr=FLAGS.lr,
          batch=FLAGS.batch,
          epochs=FLAGS.epochs,
          weight_decay=FLAGS.weight_decay,
      )
      logdir = "%s/%s" % (
          tm.__class__.__name__,
          "_".join(sorted("%s_%s" % k for k in tm.params.items())),
      )
      logdir = os.path.join(FLAGS.logdir, logdir)
      print(f'{" Model ":-^79}')
      print(tm.model)
      print(f'{" Weights ":-^79}')
      print(tm.model.vars())
      print(f'{" Config ":-^79}')
      for k, v in sorted(tm.params.items()):
          print(f"{k:32} {v}")
      print(f"Saving to {logdir}")
      print(
          f"Visualize results with:\n    tensorboard --logdir {FLAGS.logdir}"
      )
      tm.train(FLAGS.epochs, train_size, train, test, logdir)
  
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 a global function with example usage: main(argv) without return types