Not a major/urgent issue, but just FYI:
In define_optimizer_flags the optimizer flag is defined one way, while in create_optimizer_fn_from_flags the optimizer flag is defined another way.
For example, if your prefix is '' then define_optimizer_flags defines the flag '_optimizer' while create_optimizer_fn_from_flags looks for 'optimizer'. I fixed it in my own repo by using the method prefixed in both define_optimizer_flags and create_optimizer_fn_from_flags.