code: def train(
      train_dir,
      config,
      dataset_fn,
      checkpoints_to_keep=5,
      keep_checkpoint_every_n_hours=1,
      num_steps=None,
      master="",
      num_sync_workers=0,
      num_ps_tasks=0,
      task=0,
  ):
      """Train loop."""
      tf.gfile.MakeDirs(train_dir)
      is_chief = task == 0
      if is_chief:
          _trial_summary(
              config.hparams,
              config.train_examples_path or config.tfds_name,
              train_dir,
          )
      with tf.Graph().as_default():
          with tf.device(
              tf.train.replica_device_setter(
                  num_ps_tasks, merge_devices=True
              )
          ):
              model = config.model
              model.build(
                  config.hparams,
                  config.data_converter.output_depth,
                  is_training=True,
              )
  
              optimizer = model.train(
                  **_get_input_tensors(dataset_fn(), config)
              )
  
              hooks = []
              if num_sync_workers:
                  optimizer = tf.train.SyncReplicasOptimizer(
                      optimizer, num_sync_workers
                  )
                  hooks.append(
                      optimizer.make_session_run_hook(is_chief)
                  )
  
              grads, var_list = list(
                  zip(*optimizer.compute_gradients(model.loss))
              )
              global_norm = tf.global_norm(grads)
              tf.summary.scalar("global_norm", global_norm)
  
              if config.hparams.clip_mode == "value":
                  g = config.hparams.grad_clip
                  clipped_grads = [
                      tf.clip_by_value(grad, -g, g) for grad in grads
                  ]
              elif config.hparams.clip_mode == "global_norm":
                  clipped_grads = tf.cond(
                      global_norm
                      < config.hparams.grad_norm_clip_to_zero,
                      lambda: tf.clip_by_global_norm(  # pylint:disable=g-long-lambda
                          grads,
                          config.hparams.grad_clip,
                          use_norm=global_norm,
                      )[
                          0
                      ],
                      lambda: [tf.zeros(tf.shape(g)) for g in grads],
                  )
              else:
                  raise ValueError(
                      "Unknown clip_mode: {}".format(
                          config.hparams.clip_mode
                      )
                  )
              train_op = optimizer.apply_gradients(
                  list(zip(clipped_grads, var_list)),
                  global_step=model.global_step,
                  name="train_step",
              )
  
              logging_dict = {
                  "global_step": model.global_step,
                  "loss": model.loss,
              }
  
              hooks.append(
                  tf.train.LoggingTensorHook(
                      logging_dict, every_n_iter=100
                  )
              )
              if num_steps:
                  hooks.append(
                      tf.train.StopAtStepHook(last_step=num_steps)
                  )
  
              scaffold = tf.train.Scaffold(
                  saver=tf.train.Saver(
                      max_to_keep=checkpoints_to_keep,
                      keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
                  )
              )
              tf_slim.training.train(
                  train_op=train_op,
                  logdir=train_dir,
                  scaffold=scaffold,
                  hooks=hooks,
                  save_checkpoint_secs=60,
                  master=master,
                  is_chief=is_chief,
              )
  
file_path: magenta/magenta/models/music_vae/music_vae_train.py
imports: import os
  from magenta.models.music_vae import configs
  from magenta.models.music_vae import data
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 3 for calling a global function with example usage: train(train_dir, config, dataset_fn, checkpoints_to_keep, keep_checkpoint_every_n_hours, num_steps, master, num_sync_workers, num_ps_tasks, task) without return types