'Build training specific ops for the graph.'

def _build_train_op(self):

    self.lrn_rate = tf.constant(self.hps.lrn_rate, tf.float32)
    tf.summary.scalar('learning_rate', self.lrn_rate)
    trainable_variables = tf.trainable_variables()
    grads = tf.gradients(self.cost, trainable_variables)
    if (self.hps.optimizer == 'sgd'):
        optimizer = tf.train.GradientDescentOptimizer(self.lrn_rate)
    elif (self.hps.optimizer == 'mom'):
        optimizer = tf.train.MomentumOptimizer(self.lrn_rate, 0.9)
    apply_op = optimizer.apply_gradients(zip(grads, trainable_variables), global_step=self.global_step, name='train_step')
    train_ops = ([apply_op] + self._extra_train_ops)
    self.train_op = tf.group(*train_ops)
