Merge pull request #2532 from ROCm/develop-upstream-QA-rocm62-tf-keras

[develop-upstream-QA-rocm62] Add tf_keras to container
ROCm · May 9, 2024 · 9ab6773 · 9ab6773
2 parents 8a3a823 + 2f443fc
commit 9ab6773
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/devel.requirements.txt b/tensorflow/tools/tf_sig_build_dockerfiles/devel.requirements.txt
@@ -32,8 +32,9 @@ gast == 0.4.0
 # For release jobs, we will pin these on the release branch
 # Note that the CACHEBUSTER variable, set in the CI builds, will force these to
 # be the latest version.
-keras-nightly ~= 2.14.0.dev
-tb-nightly ~= 2.13.0.a
+tb-nightly ~= 2.17.0a20240220
+keras-nightly ~= 3.1.0.dev2024022003
+tf-keras-nightly ~= 2.17.0.dev2024022010
 # Test dependencies
 grpcio ~= 1.59.0 # Earliest version for Python 3.12
 portpicker ~= 1.6.0

diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/setup.python.sh b/tensorflow/tools/tf_sig_build_dockerfiles/setup.python.sh
@@ -90,8 +90,10 @@ python3 -m pip install -U setuptools
 
 if [[ $3 ]]; then
     echo "Runtime mode"
+    # Install tf_keras
+    python3 -m pip install --no-cache-dir --no-deps tf-keras-nightly==2.17.0.dev2024022010
 else
-    echo "Install Requirements"
+    echo "Install Dev Requirements"
     # Disable the cache dir to save image space, and install packages
     python3 -m pip install --no-cache-dir -r $REQUIREMENTS -U
 fi