tfm.optimization.register_optimizer_cls

Register customize optimizer cls.

The user will still need to subclass data classes in configs.optimization_config to be used with OptimizerFactory.

key A string to that the optimizer_config_cls is registered with.
optimizer_config_cls A class which inherits tf.keras.optimizers.Optimizer.
use_legacy_optimizer A boolean that indicates if using legacy optimizers.