diff --git a/tools/docker/build_wheel.Dockerfile b/tools/docker/build_wheel.Dockerfile index 9542e1365a..ef64befbd2 100644 --- a/tools/docker/build_wheel.Dockerfile +++ b/tools/docker/build_wheel.Dockerfile @@ -1,6 +1,6 @@ #syntax=docker/dockerfile:1.1.5-experimental ARG PY_VERSION -FROM tensorflow/build:2.15-python$PY_VERSION as base_install +FROM tensorflow/build:2.16-python$PY_VERSION as base_install ENV TF_NEED_CUDA="1" ARG PY_VERSION