aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/estimator
diff options
context:
space:
mode:
authorGravatar Yan Facai (颜发才) <facai.yan@gmail.com>2018-08-21 13:06:02 +0800
committerGravatar Yan Facai (颜发才) <facai.yan@gmail.com>2018-08-21 13:35:31 +0800
commit5630efcca924563b549a788b4b5ec93fea91e559 (patch)
tree402a0735cbeba8427040f655a40043fa9b414a54 /tensorflow/contrib/estimator
parent52d637e604dacd3bff836a27bd991f95966226e8 (diff)
CLN: revise according to comments
Diffstat (limited to 'tensorflow/contrib/estimator')
-rw-r--r--tensorflow/contrib/estimator/python/estimator/boosted_trees.py5
1 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/contrib/estimator/python/estimator/boosted_trees.py b/tensorflow/contrib/estimator/python/estimator/boosted_trees.py
index 7ed77bcce6..e6bdc97fe5 100644
--- a/tensorflow/contrib/estimator/python/estimator/boosted_trees.py
+++ b/tensorflow/contrib/estimator/python/estimator/boosted_trees.py
@@ -33,7 +33,7 @@ def _validate_input_fn_and_repeat_dataset(train_input_fn):
return _input_fn
-class _BoostedTreesEstimator(estimator.Estimator):
+class _BoostedTreesEstimator(canned_boosted_trees._BoostedTrees): # pylint: disable=protected-access
"""An Estimator for Tensorflow Boosted Trees models."""
def __init__(self,
@@ -115,7 +115,8 @@ class _BoostedTreesEstimator(estimator.Estimator):
config=config)
super(_BoostedTreesEstimator, self).__init__(
- model_fn=_model_fn, model_dir=model_dir, config=config)
+ model_fn=_model_fn, model_dir=model_dir, config=config,
+ feature_columns=feature_columns)
# pylint:enable=protected-access