diff --git a/.dockerignore b/.dockerignore
deleted file mode 100644
index 6b8710a..0000000
--- a/.dockerignore
+++ /dev/null
@@ -1 +0,0 @@
-.git
diff --git a/pip_pkg_scripts/README.md b/pip_pkg_scripts/README.md
index bf6e32b..ec711b0 100644
--- a/pip_pkg_scripts/README.md
+++ b/pip_pkg_scripts/README.md
@@ -2,19 +2,20 @@
 
 ## Build a pip package locally using the docker container
 
-```bash
+```
 docker build --tag=open_dataset_pip -f pip_pkg_scripts/build.Dockerfile .
-mkdir /tmp/artifacts
-docker run --mount type=bind,source=/tmp/artifacts,target=/tmp/artifacts -e "PYTHON_VERSION=3" -e "PYTHON_MINOR_VERSION=8" -e "PIP_MANYLINUX2010=1" -e "TF_VERSION=2.3.0" open_dataset_pip
+mkdir /tmp/pip_pkg_build
+docker run --mount type=bind,source=/tmp/pip_pkg_build,target=/tmp/pip_pkg_build -e "GITHUB_BRANCH=master" -e "PYTHON_VERSION=3" -e "PYTHON_MINOR_VERSION=8" -e "PIP_MANYLINUX2010=1" -e "TF_VERSION=2.3.0" open_dataset_pip
 ```
 This command will execute the `build.sh` inside the container, which clones the
 github repository, builds the library and outputs `.whl` packages under
-`/tmp/artifacts/`
+`/tmp/pip_pkg_build/`
 
 ## Build a pip package locally without docker container
 First follow quick start to install all the depdencies such as bazel. Then run
 
-```bash
+```
+export GITHUB_BRANCH=master
 export PYTHON_VERSION=3
 export PYTHON_MINOR_VERSION=7
 export PIP_MANYLINUX2010=0
diff --git a/pip_pkg_scripts/build.Dockerfile b/pip_pkg_scripts/build.Dockerfile
index a2f48f5..d98142c 100644
--- a/pip_pkg_scripts/build.Dockerfile
+++ b/pip_pkg_scripts/build.Dockerfile
@@ -1,5 +1,6 @@
 FROM tensorflow/tensorflow:custom-op-ubuntu16
 
+ENV GITHUB_BRANCH="master"
 ENV PYTHON_VERSION="3"
 ENV PYTHON_MINOR_VERSION=""
 ENV PIP_MANYLINUX2010="1"
@@ -25,11 +26,11 @@ RUN for python in python3.5 python3.6 python3.7 python3.8; do \
       $python -m pip install --upgrade grpcio>=1.24.3; \
     done
 
-VOLUME /tmp/artifacts
-COPY . /tmp/repo
-WORKDIR /tmp/repo
+COPY pip_pkg_scripts/build.sh /
 
-ENTRYPOINT ["pip_pkg_scripts/build.sh"]
+VOLUME /tmp/pip_pkg_build
+
+ENTRYPOINT ["/build.sh"]
 
 # The default parameters for the build.sh
 CMD []
diff --git a/pip_pkg_scripts/build.sh b/pip_pkg_scripts/build.sh
index 29fbadf..1274347 100755
--- a/pip_pkg_scripts/build.sh
+++ b/pip_pkg_scripts/build.sh
@@ -22,6 +22,7 @@
 set -e -x
 
 # Override the following env variables if necessary.
+export GITHUB_BRANCH="${GITHUB_BRANCH:-master}"
 export PYTHON_VERSION="${PYTHON_VERSION:-3}"
 export PYTHON_MINOR_VERSION="${PYTHON_MINOR_VERSION}"
 export PIP_MANYLINUX2010="${PIP_MANYLINUX2010:-0}"
@@ -45,13 +46,19 @@ ${PIP} install tensorflow-gpu=="${TF_VERSION}" --user
 # bazel.
 python3 -m pip install --upgrade tensorflow-gpu=="${TF_VERSION}" --user
 
+rm -rf waymo-od || true
+git clone https://github.com/waymo-research/waymo-open-dataset.git waymo-od
+cd waymo-od
+
+git checkout remotes/origin/${GITHUB_BRANCH}
+
 ./configure.sh
 
 bazel clean
 bazel build ...
 bazel test ...
 
-DST_DIR="/tmp/artifacts"
+DST_DIR="/tmp/pip_pkg_build"
 rm -rf "$DST_DIR" || true
 ./pip_pkg_scripts/build_pip_pkg.sh "$DST_DIR" ${PYTHON_VERSION}
 # Comment the following line if you run this outside of the container.