diff --git a/requirements.txt b/requirements.txt index c28ccabba..74fce87e5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,7 +8,7 @@ scipy # with both new and old protobuf stubs. This is needed to resolve # Keras-TensorFlow circular dependency issue, when one of them gets a dependency # incompatible with another one (protobuf in this specific case). -protobuf==3.20.3 +protobuf==4.21.6 tf-nightly portpicker pyyaml @@ -16,4 +16,16 @@ Pillow numpy >= 1.26.0, < 2.1.0 # Sync with the numpy version used in TF black==22.3.0 isort==5.10.1 -flake8==4.0.1 \ No newline at end of file +flake8==4.0.1 +# NVIDIA CUDA dependencies +nvidia-cublas-cu12 == 12.5.3.2 +nvidia-cuda-cupti-cu12 == 12.5.82 +nvidia-cuda-nvrtc-cu12 == 12.5.82 +nvidia-cuda-runtime-cu12 == 12.5.82 +nvidia-cudnn-cu12 == 9.3.0.75 +nvidia-cufft-cu12 == 11.2.3.61 +nvidia-curand-cu12 == 10.3.6.82 +nvidia-cusolver-cu12 == 11.6.3.83 +nvidia-cusparse-cu12 == 12.5.1.3 +nvidia-nccl-cu12 == 2.25.1 +nvidia-nvjitlink-cu12 == 12.5.82 diff --git a/tf_keras/kokoro/github/ubuntu/gpu/build.sh b/tf_keras/kokoro/github/ubuntu/gpu/build.sh index 0c9bfa4ce..386d18e15 100644 --- a/tf_keras/kokoro/github/ubuntu/gpu/build.sh +++ b/tf_keras/kokoro/github/ubuntu/gpu/build.sh @@ -40,9 +40,6 @@ pip install -r requirements.txt # Update this once we switch to `tf_keras-nightly` in TensorFlow. pip uninstall -y keras-nightly -# LD Library Path needs to be same as TensorFlow Ubuntu Docker build - -# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/tf_sig_build_dockerfiles/ -export LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib" CUDA_TOOLKIT_PATH="/usr/local/cuda-11.8" TF_CUDA_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda11.8-cudnn8.6-tensorrt8.4_config_cuda" TF_CUDA_COMPUTE_CAPABILITIES="sm_35,sm_50,sm_60,sm_70,sm_75,compute_80"