def get_standard_adamw_optimizer(
    learning_rate_fn: Callable[[int], float],
    config: ml_collections.ConfigDict) -> optax.GradientTransformation:
  """Constructs the optimizer from the given HParams."""
  weight_decay = config.optimizer.get('weight_decay', 0.0)
  tx = optax.adamw(  # pytype: disable=wrong-arg-types  # numpy-scalars
      learning_rate=learning_rate_fn,
      b1=config.optimizer.beta1,
      b2=config.optimizer.beta2,
      weight_decay=weight_decay)
  return tx
Constructs the optimizer from the given HParams.
videogvt/train_lib/maskgvt_train_lib.py