from tensorflow.python.ops import array_ops

        array_ops.reshape(centered_bias, [-1]))
    return centered_bias

  def _centered_bias_step(self, targets, features):
    centered_bias = ops.get_collection(self._centered_bias_weight_collection)
    batch_size = array_ops.shape(targets)[0]
    logits = array_ops.reshape(
        array_ops.tile(centered_bias[0], [batch_size]),
        [batch_size, self._target_column.num_label_columns])
    loss = self._target_column.loss(logits, targets, features)
    # Learn central bias by an optimizer. 0.1 is a convervative lr for a single
    # variable.
    return training.AdagradOptimizer(0.1).minimize(loss, var_list=centered_bias)

