code: def train(
      config_name,
      ckpt_dir,
      num_steps: int,
      auto_encoder_ckpt_dir,
      batch_size,
      crop_size,
      lpips_weight_path,
      create_image_summaries,
      tfds_arguments: model.TFDSArguments,
  ):
      """Train the model."""
      config = configs.get_config(config_name)
      hific = model.HiFiC(
          config,
          helpers.ModelMode.TRAINING,
          lpips_weight_path,
          auto_encoder_ckpt_dir,
          create_image_summaries,
      )
  
      dataset = hific.build_input(
          batch_size, crop_size, tfds_arguments=tfds_arguments
      )
      iterator = tf.data.make_one_shot_iterator(dataset)
      get_next = iterator.get_next()
  
      hific.build_model(**get_next)
      train_op = hific.train_op
  
      hooks = hific.hooks + [
          tf.train.StopAtStepHook(last_step=num_steps)
      ]
      global_step = tf.train.get_or_create_global_step()
      tf.logging.info(
          f"\nStarting MonitoredTrainingSession at {ckpt_dir}\n"
      )
  
      with tf.train.MonitoredTrainingSession(
          checkpoint_dir=ckpt_dir,
          save_checkpoint_steps=SAVE_CHECKPOINT_STEPS,
          hooks=hooks,
      ) as sess:
          if auto_encoder_ckpt_dir:
              hific.restore_autoencoder(sess)
          tf.logging.info("Session setup, starting training...")
          while True:
              if sess.should_stop():
                  break
              global_step_np, _ = sess.run([global_step, train_op])
              if global_step_np == 0:
                  tf.logging.info("First iteration passed.")
              if global_step_np > 1 and global_step_np % 100 == 1:
                  tf.logging.info(f"Iteration {global_step_np}")
      tf.logging.info("Training session closed.")
  
file_path: compression/models/hific/train.py
imports: import argparse
  import sys
  import tensorflow.compat.v1
  from hific import configs
  from hific import helpers
  from hific import model
label_desc: Use image node_id 1 for calling a global function with example usage: train(config_name, ckpt_dir, num_steps, auto_encoder_ckpt_dir, batch_size, crop_size, lpips_weight_path, create_image_summaries, tfds_arguments) without return types