From 601488fd4c1468ae7872e132e0f1c9843df54182 Mon Sep 17 00:00:00 2001
From: Ramesh Sampath <1437573+sampathweb@users.noreply.github.com>
Date: Mon, 6 Nov 2023 17:55:08 -0600
Subject: [PATCH] Use TENSORFLOW_VERSION when available during pip_build script
 (#18739)

---
 pip_build.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pip_build.py b/pip_build.py
index 6c09c1ccb7a..d7604319439 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -434,6 +434,9 @@ def test_wheel(wheel_path, expected_version, requirements_path):
         "virtualenv kenv\n"
         f"source {os.path.join('kenv', 'bin', 'activate')}\n"
         f"pip3 install -r {requirements_path}\n"
+        "pip3 uninstall -y tensorflow tf-nightly\n"
+        'pip3 install "${TENSORFLOW_VERSION:-tf-nightly}"\n'
+        "pip3 uninstall -y keras keras-nightly\n"
         f"pip3 install {wheel_path} --force-reinstall\n"
         f"python3 -c 'import keras;{checks};print(keras.__version__)'\n"
         f"python3 -c 'import tensorflow as tf;tf.compat.v1.layers.Dense'\n"