Handle missing contrib import

tensorflow · Sep 19, 2019 · 4e49cc8 · 4e49cc8
1 parent 269a406
commit 4e49cc8
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/adanet/core/estimator_distributed_test_runner.py b/adanet/core/estimator_distributed_test_runner.py
@@ -50,9 +50,14 @@
 # pylint: disable=g-direct-tensorflow-import
 
 # Contrib
-from tensorflow.contrib.boosted_trees.estimator_batch.estimator import CoreGradientBoostedDecisionTreeEstimator
-from tensorflow.contrib.boosted_trees.proto import learner_pb2
-from tensorflow.contrib.boosted_trees.python.utils import losses as bt_losses
+try:
+  from tensorflow.contrib.boosted_trees.estimator_batch.estimator import CoreGradientBoostedDecisionTreeEstimator
+  from tensorflow.contrib.boosted_trees.proto import learner_pb2
+  from tensorflow.contrib.boosted_trees.python.utils import losses as bt_losses
+except ImportError:
+  # Not much we can do here except skip the test.
+  pass
+
 from tensorflow.python.ops import partitioned_variables
 from tensorflow.python.training import session_manager as session_manager_lib
 from tensorflow_estimator.python.estimator.canned import head as head_lib