diff options
Diffstat (limited to 'tensorflow/python/estimator/canned/boosted_trees.py')
| -rw-r--r-- | tensorflow/python/estimator/canned/boosted_trees.py | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/tensorflow/python/estimator/canned/boosted_trees.py b/tensorflow/python/estimator/canned/boosted_trees.py index 4e6010a162..6b54f51ca6 100644 --- a/tensorflow/python/estimator/canned/boosted_trees.py +++ b/tensorflow/python/estimator/canned/boosted_trees.py @@ -39,7 +39,7 @@ from tensorflow.python.summary import summary from tensorflow.python.training import distribute as distribute_lib from tensorflow.python.training import session_run_hook from tensorflow.python.training import training_util -from tensorflow.python.util.tf_export import tf_export +from tensorflow.python.util.tf_export import estimator_export # TODO(nponomareva): Reveal pruning params here. _TreeHParams = collections.namedtuple('TreeHParams', [ @@ -712,7 +712,7 @@ def _create_regression_head(label_dimension, weight_column=None): # pylint: enable=protected-access -@tf_export('estimator.BoostedTreesClassifier') +@estimator_export('estimator.BoostedTreesClassifier') class BoostedTreesClassifier(estimator.Estimator): """A Classifier for Tensorflow Boosted Trees models.""" @@ -830,7 +830,7 @@ class BoostedTreesClassifier(estimator.Estimator): model_fn=_model_fn, model_dir=model_dir, config=config) -@tf_export('estimator.BoostedTreesRegressor') +@estimator_export('estimator.BoostedTreesRegressor') class BoostedTreesRegressor(estimator.Estimator): """A Regressor for Tensorflow Boosted Trees models.""" |