diff --git a/tools/docker/build_wheel.Dockerfile b/tools/docker/build_wheel.Dockerfile
index 9a70817178..2052fab043 100644
--- a/tools/docker/build_wheel.Dockerfile
+++ b/tools/docker/build_wheel.Dockerfile
@@ -1,33 +1,60 @@
 #syntax=docker/dockerfile:1.1.5-experimental
 ARG TF_VERSION
-FROM seanpmorgan/tensorflow:2.1.0-custom-op-gpu-ubuntu16-minimal as make_wheel
+FROM seanpmorgan/tensorflow:2.1.0-custom-op-gpu-ubuntu16-minimal as base_install
 ENV TF_NEED_CUDA="1"
 
 RUN apt-get update && apt-get install patchelf
 
 ARG PY_VERSION
-RUN python$PY_VERSION -m pip install --upgrade pip setuptools auditwheel==2.0.0
+RUN ln -sf $(which python$PY_VERSION) /usr/bin/python3
+
+RUN python3 -m pip install --upgrade pip setuptools auditwheel==2.0.0
 
 COPY tools/install_deps/ /install_deps
 ARG TF_VERSION
-RUN python$PY_VERSION -m pip install \
+RUN python3 -m pip install \
         tensorflow==$TF_VERSION \
         -r /install_deps/pytest.txt
 
 COPY requirements.txt .
-RUN python$PY_VERSION -m pip install -r requirements.txt
+RUN python3 -m pip install -r requirements.txt
+
+# Fix presented in
+# https://stackoverflow.com/questions/44967202/pip-is-showing-error-lsb-release-a-returned-non-zero-exit-status-1/44967506
+RUN echo "#! /usr/bin/python2.7" >> /usr/bin/lsb_release2
+RUN cat /usr/bin/lsb_release >> /usr/bin/lsb_release2
+RUN mv /usr/bin/lsb_release2 /usr/bin/lsb_release
 
 COPY ./ /addons
 WORKDIR /addons
+
+# -------------------------------------------------------------------
+FROM base_install as tfa_gpu_tests
+CMD ["bash", "tools/testing/build_and_run_tests.sh"]
+
+# -------------------------------------------------------------------
+FROM base_install as make_wheel
 ARG NIGHTLY_FLAG
 ARG NIGHTLY_TIME
 RUN --mount=type=cache,id=cache_bazel,target=/root/.cache/bazel \
-    bash tools/releases/release_linux.sh $PY_VERSION $NIGHTLY_FLAG
+    bash tools/testing/build_and_run_tests.sh && \
+    bazel clean --expunge && \
+    bazel build \
+        -c opt \
+        --noshow_progress \
+        --noshow_loading_progress \
+        --verbose_failures \
+        --test_output=errors \
+        --crosstool_top=//build_deps/toolchains/gcc7_manylinux2010-nvcc-cuda10.1:toolchain \
+        build_pip_pkg && \
+    # Package Whl
+    bazel-bin/build_pip_pkg artifacts $NIGHTLY_FLAG
 
 RUN bash tools/releases/tf_auditwheel_patch.sh
 RUN auditwheel repair --plat manylinux2010_x86_64 artifacts/*.whl
 RUN ls -al wheelhouse/
 
+# -------------------------------------------------------------------
 FROM scratch as output
 
 COPY --from=make_wheel /addons/wheelhouse/ .
diff --git a/tools/docker/gpu_tests.Dockerfile b/tools/docker/gpu_tests.Dockerfile
deleted file mode 100644
index 6f1f52045f..0000000000
--- a/tools/docker/gpu_tests.Dockerfile
+++ /dev/null
@@ -1,15 +0,0 @@
-FROM tensorflow/tensorflow:2.1.0-custom-op-gpu-ubuntu16
-ENV TF_NEED_CUDA="1"
-
-RUN python3 -m pip install --upgrade pip setuptools auditwheel==2.0.0
-
-COPY tools/install_deps/tensorflow.txt ./
-RUN python3 -m pip install -r tensorflow.txt
-
-COPY requirements.txt ./
-RUN python3 -m pip install -r requirements.txt
-
-
-COPY ./ /addons
-WORKDIR addons
-CMD ["bash", "tools/testing/build_and_run_tests.sh"]
diff --git a/tools/releases/release_linux.sh b/tools/releases/release_linux.sh
deleted file mode 100644
index 0b9b93d2e4..0000000000
--- a/tools/releases/release_linux.sh
+++ /dev/null
@@ -1,54 +0,0 @@
-#!/usr/bin/env bash
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-set -e -x
-
-if [[ $2 == "--nightly" ]]; then
-    PKG_OPS="--nightly"
-elif [[ -n "$2" ]]; then
-    echo "Found unsupported args: $@"
-    exit 1
-fi
-
-# Configs
-export TF_NEED_CUDA="1"
-
-# Fix presented in
-# https://stackoverflow.com/questions/44967202/pip-is-showing-error-lsb-release-a-returned-non-zero-exit-status-1/44967506
-echo "#! /usr/bin/python2.7" >> /usr/bin/lsb_release2
-cat /usr/bin/lsb_release >> /usr/bin/lsb_release2
-mv /usr/bin/lsb_release2 /usr/bin/lsb_release
-
-ln -sf $(which python$1) /usr/bin/python3
-python3 -m pip install --upgrade pip
-python3 -m pip install --upgrade setuptools
-
-#Link TF dependency
-python3 --version
-bash tools/testing/build_and_run_tests.sh
-bazel clean --expunge
-
-# Build
-bazel build \
-  -c opt \
-  --noshow_progress \
-  --noshow_loading_progress \
-  --verbose_failures \
-  --test_output=errors \
-  --crosstool_top=//build_deps/toolchains/gcc7_manylinux2010-nvcc-cuda10.1:toolchain \
-  build_pip_pkg
-
-# Package Whl
-bazel-bin/build_pip_pkg artifacts ${PKG_OPS}
diff --git a/tools/run_gpu_tests.sh b/tools/run_gpu_tests.sh
index bea1e96861..cf81ae3925 100644
--- a/tools/run_gpu_tests.sh
+++ b/tools/run_gpu_tests.sh
@@ -3,5 +3,10 @@
 set -x -e
 
 export DOCKER_BUILDKIT=1
-docker build -f tools/docker/gpu_tests.Dockerfile -t tfa_gpu_tests ./
-docker run --rm -t --runtime=nvidia tfa_gpu_tests
+docker build \
+       -f tools/docker/build_wheel.Dockerfile \
+       --target tfa_gpu_tests \
+       --build-arg TF_VERSION=2.1.0 \
+       --build-arg PY_VERSION=3.5 \
+       -t tfa_gpu_tests ./
+docker run --rm -t -v cache_bazel:/root/.cache/bazel --runtime=nvidia tfa_gpu_tests