From 6beb8734fbd7b247f96b869ab547882c6c004130 Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Thu, 13 Dec 2018 00:02:41 -0800
Subject: [PATCH] Update TF 2.0 package to require Estimator 2.0 package.

PiperOrigin-RevId: 225321405
---
 tensorflow/tools/pip_package/setup.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index 85c913f1588..f4b27fed731 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -89,6 +89,8 @@ if 'tf_nightly' in project_name:
       REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.13.0a0, < 1.14.0a0'
     if 'tensorflow_estimator' in pkg:
       REQUIRED_PACKAGES[i] = 'tf-estimator-nightly'
+    if 'tensorflow_estimator' in pkg and '2.0' in project_name:
+      REQUIRED_PACKAGES[i] = 'tensorflow-estimator-2.0-preview'
 
 # weakref.finalize and enum were introduced in Python 3.4
 if sys.version_info < (3, 4):