from tensorflow.python.estimator.canned import head as head_lib



class CoreGradientBoostedDecisionTreeEstimator(test_util.TensorFlowTestCase):

  def testTrainEvaluateInferDoesNotThrowError(self):
    head_fn = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss(
        loss_reduction=losses.Reduction.SUM_OVER_NONZERO_WEIGHTS)

    learner_config = learner_pb2.LearnerConfig()
    learner_config.num_classes = 2
