def get_optimizer(config: ml_collections.ConfigDict):
  """Constructs the VQGAN optimizer from the given HParams."""

  assert (
      'lr_configs' in config and 'learning_rate_schedule' in config.lr_configs
  ), 'Only support learning_rate_schedule'

  lr_config = config['lr_configs']
  g_lr_config = copy.deepcopy(lr_config).unlock()
  g_lr_config.base_learning_rate = config.optimizer.g_lr
  g_lr_config.lock()

  d_lr_config = copy.deepcopy(lr_config).unlock()
  d_lr_config.base_learning_rate = config.optimizer.d_lr
  d_lr_config.lock()

  lr_meta_fn = lr_schedules.lr_fn_dict[g_lr_config['learning_rate_schedule']]
  g_lr_fn = lr_meta_fn(g_lr_config)
  d_lr_fn = lr_meta_fn(d_lr_config)

  tx_dict = vqgan_train_lib.get_optimizer(g_lr_fn, d_lr_fn, config=config)
  return tx_dict, g_lr_fn, d_lr_fn
Constructs the VQGAN optimizer from the given HParams.
videogvt/trainers/vqgan_trainer.py