code: def main(argv):
      assert (
          len(argv) == 1
      ), "Please specify arguments via flags. Use --help for instructions"
  
      assert (
          getattr(elegy.nets.resnet, FLAGS.model, None) is not None
      ), f"{FLAGS.model} is not defined in elegy.nets.resnet"
  
      assert not os.path.exists(
          FLAGS.output_dir
      ), "Output directory already exists. Delete manually or specify a new one."
      os.makedirs(FLAGS.output_dir)
  
      # dataset
      dataset_builder = tfds.builder(FLAGS.dataset)
      ds_train = input_pipeline.create_split(
          dataset_builder,
          batch_size=FLAGS.batch_size,
          image_size=FLAGS.image_size,
          dtype=FLAGS.dtype,
          train=True,
          cache=FLAGS.cache,
      )
      ds_valid = input_pipeline.create_split(
          dataset_builder,
          batch_size=FLAGS.batch_size,
          image_size=FLAGS.image_size,
          dtype=FLAGS.dtype,
          train=False,
          cache=FLAGS.cache,
      )
      N_BATCHES_TRAIN = (
          dataset_builder.info.splits["train"].num_examples
          // FLAGS.batch_size
      )
      N_BATCHES_VALID = (
          dataset_builder.info.splits["validation"].num_examples
          // FLAGS.batch_size
      )
  
      # generator that converts tfds dataset batches to jax arrays
      def tfds2jax_generator(tf_ds):
          for batch in tf_ds:
              yield jnp.asarray(
                  batch["image"], dtype=FLAGS.dtype
              ), jax.device_put(jnp.asarray(batch["label"]))
  
      # model and optimizer definition
      def build_optimizer(
          lr,
          momentum,
          steps_per_epoch,
          n_epochs,
          nesterov,
          warmup_epochs=5,
      ):
          cosine_schedule = optax.cosine_decay_schedule(
              1, decay_steps=n_epochs * steps_per_epoch, alpha=1e-10
          )
          warmup_schedule = optax.polynomial_schedule(
              init_value=0.0,
              end_value=1.0,
              power=1,
              transition_steps=warmup_epochs * steps_per_epoch,
          )
          schedule = lambda x: jnp.minimum(
              cosine_schedule(x), warmup_schedule(x)
          )
          optimizer = optax.sgd(lr, momentum, nesterov=nesterov)
          optimizer = optax.chain(
              optimizer, optax.scale_by_schedule(schedule)
          )
          return optimizer
  
      module = getattr(elegy.nets.resnet, FLAGS.model)(
          dtype=FLAGS.dtype
      )
      model = elegy.Model(
          module,
          loss=[
              elegy.losses.Crossentropy(
                  from_logits=True, weight=FLAGS.loss_scale
              ),
              elegy.regularizers.L2(
                  FLAGS.L2_reg / 2 * FLAGS.loss_scale
              ),
          ],
          metrics=elegy.metrics.Accuracy(),
          optimizer=build_optimizer(
              FLAGS.base_lr / FLAGS.loss_scale,
              FLAGS.momentum,
              N_BATCHES_TRAIN,
              FLAGS.epochs,
              FLAGS.nesterov,
          ),
      )
  
      # training
      model.fit(
          inputs=tfds2jax_generator(ds_train),
          validation_data=tfds2jax_generator(ds_valid),
          epochs=FLAGS.epochs,
          verbose=2,
          steps_per_epoch=N_BATCHES_TRAIN,
          validation_steps=N_BATCHES_VALID,
          callbacks=[
              elegy.callbacks.ModelCheckpoint(
                  FLAGS.output_dir, save_best_only=True
              ),
              elegy.callbacks.TerminateOnNaN(),
              elegy.callbacks.TensorBoard(logdir=FLAGS.output_dir),
          ],
      )
  
file_path: elegy/examples/need-fixing/imagenet/resnet_imagenet.py
imports: import os
  import jax
  import jax.numpy
  from absl import app, flags
  import input_pipeline
  import optax
  import tensorflow_datasets
  import elegy
label_desc: Use image node_id 1 for calling a global function with example usage: main(argv) and returns: optimizer