def main(rng: jnp.ndarray, config: ml_collections.ConfigDict, workdir: str,
         writer: metric_writers.MetricWriter):
  """Main function for the VideoGVT project."""
  rng, global_seed_rng, local_rng = jax.random.split(rng, 3)
  # Make sure each host uses a different RNG for the training data.
  is_train = not FLAGS.eval_while_training and not config.get('eval_only')

  local_rng = jax.random.fold_in(local_rng, jax.process_index())
  data_rng, local_seed_rng = jax.random.split(local_rng)
  train_utils.seed_everything(int(global_seed_rng[0]), int(local_seed_rng[0]))
  trainer = get_trainer(config.model_class, is_train=is_train)
  config = update_config_for_evaluation(config)

  dataset = train_utils.get_dataset(
      config,
      data_rng,
      is_train,
      dataset_service_address=FLAGS.dataset_service_address)
  logging.info('Trainer %s loaded', config.model_class)
  ret = trainer(
      rng=rng, config=config, dataset=dataset, workdir=workdir, writer=writer)
  # Wait until all processes are done before exiting.
  jax.pmap(jax.random.PRNGKey)(jnp.arange(
      jax.local_device_count())).block_until_ready()
  return ret
Main function for the VideoGVT project.
videogvt/main.py