import tensorflow as tf

            if self.mode == tf.estimator.ModeKeys.EVAL:

                return tf.estimator.EstimatorSpec(
                    self.mode, loss=loss, eval_metric_ops=metrics
                )

            elif self.mode == tf.estimator.ModeKeys.TRAIN:

                optimizer_params = self.params.get("optimizer_params", {})
                global_step = tf.train.get_or_create_global_step()

                # apply learning rate decay if it's setup already.
                lr_decay_params = optimizer_params.pop("learning_rate_exp_decay", {})

                # learning_rate = tf.train.exponential_decay(
                #     self.params["learning_rate"],
                #     global_step,
                #     decay_steps=self.params["lr_decay_steps"],
