In the list of features for nn_adam, you are using this feature, but if I understand correctly, this feature is always 1. https://github.com/google/learned_optimization/blob/242e2180f03e8dfa72a08cd0d52126951b471909/learned_optimization/learned_optimizers/nn_adam.py#L305