Build training specific ops for the graph.

def build_train_op_from_cost(self, cost):
    """Build training specific ops for the graph."""
    self.lrn_rate = tf.constant(self.hps.lrn_rate, tf.float32,
                                name='learning_rate')
    self.momentum = tf.constant(self.hps.momentum, tf.float32,
                                name='momentum')

    trainable_variables = tf.trainable_variables()
    grads = tf.gradients(cost, trainable_variables)
    devs = {v.device for v in trainable_variables}
    assert len(devs) == 1, ('There should be no trainable variables'
                            ' on any device other than the last GPU.')

    optimizer = tf.train.MomentumOptimizer(self.lrn_rate, self.momentum)

    gv_pairs = zip(grads, trainable_variables)
    gv_pairs = [gv for gv in gv_pairs if gv[0] is not None]
    devs = {gv[1].device for gv in gv_pairs}
    assert len(devs) == 1, ('There should be no gradients wrt'
                            ' vars on other GPUs.')

    apply_op = optimizer.apply_gradients(
        gv_pairs,
        global_step=self.global_step, name='train_step')

    train_ops = [apply_op]
    train_op = tf.group(*train_ops)
    return train_op