def get_optimizer(
    g_learning_rate_fn: Callable[[int], float],
    d_learning_rate_fn: Callable[[int], float],
    config: ml_collections.ConfigDict,
) -> Dict[str, optax.GradientTransformation]:
  """Constructs the optimizer from the given HParams."""
  weight_decay = config.optimizer.get('weight_decay', 0.0)
  g_tx = optax.adamw(  # pytype: disable=wrong-arg-types  # numpy-scalars
      learning_rate=g_learning_rate_fn,
      b1=config.optimizer.beta1,
      b2=config.optimizer.beta2,
      weight_decay=weight_decay)
  d_tx = optax.adamw(  # pytype: disable=wrong-arg-types  # numpy-scalars
      learning_rate=d_learning_rate_fn,
      b1=config.optimizer.beta1,
      b2=config.optimizer.beta2,
      weight_decay=weight_decay)
  if config.get('vqgan.finetune_decoder'):
    logging.info('Finetuning the VQVAE decoder only.')
    param_labels = {'decoder': 'adamw', 'encoder': 'zero', 'quantizer': 'zero'}
    g_tx = optax.multi_transform(
        {'adamw': g_tx,
         'zero': zero_grads()},
        param_labels)
  tx_dict = dict(generator=g_tx, discriminator=d_tx)
  return tx_dict
Constructs the optimizer from the given HParams.
videogvt/train_lib/vqgan_train_lib.py