Fix CI failure related to SGD optimizer. #1

PiperOrigin-RevId: 276556571
tensorflow · Oct 28, 2019 · 001d298 · 001d298
1 parent d582022
commit 001d298
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/adanet/autoensemble/estimator_v2_test.py b/adanet/autoensemble/estimator_v2_test.py
@@ -58,14 +58,14 @@ def tearDown(self):
           "candidate_pool":
               lambda head, feature_columns, optimizer: lambda config: {
                   "dnn":
-                      tf.estimator.DNNEstimator(
+                      tf.compat.v2.estimator.DNNEstimator(
                           head=head,
                           feature_columns=feature_columns,
                           optimizer=optimizer,
                           hidden_units=[3],
                           config=config),
                   "linear":
-                      tf.estimator.LinearEstimator(
+                      tf.compat.v2.estimator.LinearEstimator(
                           head=head,
                           feature_columns=feature_columns,
                           optimizer=optimizer,

diff --git a/adanet/autoensemble/keras_test.py b/adanet/autoensemble/keras_test.py
@@ -35,13 +35,13 @@ class KerasTest(parameterized.TestCase, tf.test.TestCase):
           "candidate_pool":
               lambda head, feature_columns, optimizer: {
                   "dnn":
-                      tf.estimator.DNNEstimator(
+                      tf.compat.v2.estimator.DNNEstimator(
                           head=head,
                           feature_columns=feature_columns,
                           optimizer=optimizer,
                           hidden_units=[3]),
                   "linear":
-                      tf.estimator.LinearEstimator(
+                      tf.compat.v2.estimator.LinearEstimator(
                           head=head,
                           feature_columns=feature_columns,
                           optimizer=optimizer),