code: def train(args):
      """Instantiates and trains the model."""
      if args.precision_policy:
          tf.keras.mixed_precision.set_global_policy(
              args.precision_policy
          )
      if args.check_numerics:
          tf.debugging.enable_check_numerics()
  
      model = MS2020Model(
          args.lmbda,
          args.num_filters,
          args.latent_depth,
          args.hyperprior_depth,
          args.num_slices,
          args.max_support_slices,
          args.num_scales,
          args.scale_min,
          args.scale_max,
      )
      model.compile(
          optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
      )
  
      if args.train_glob:
          train_dataset = get_custom_dataset("train", args)
          validation_dataset = get_custom_dataset("validation", args)
      else:
          train_dataset = get_dataset("clic", "train", args)
          validation_dataset = get_dataset("clic", "validation", args)
      validation_dataset = validation_dataset.take(
          args.max_validation_steps
      )
  
      model.fit(
          train_dataset.prefetch(8),
          epochs=args.epochs,
          steps_per_epoch=args.steps_per_epoch,
          validation_data=validation_dataset.cache(),
          validation_freq=1,
          callbacks=[
              tf.keras.callbacks.TerminateOnNaN(),
              tf.keras.callbacks.TensorBoard(
                  log_dir=args.train_path,
                  histogram_freq=1,
                  update_freq="epoch",
              ),
              tf.keras.callbacks.BackupAndRestore(args.train_path),
          ],
          verbose=int(args.verbose),
      )
      model.save(args.model_path)
  
imports: import argparse
  import functools
  import glob
  import sys
  from absl import app
  from absl.flags import argparse_flags
  import tensorflow
  import tensorflow_compression
  import tensorflow_datasets
label_desc: Use image node_id 7 for calling a global function with example usage: train(args) without return types