diff --git a/tensorflow/lite/tools/pip_package/Dockerfile b/tensorflow/lite/tools/pip_package/Dockerfile
new file mode 100644
index 00000000000..e4d47d0fa0d
--- /dev/null
+++ b/tensorflow/lite/tools/pip_package/Dockerfile
@@ -0,0 +1,36 @@
+ARG IMAGE
+FROM ${IMAGE}
+
+COPY update_sources.sh /
+RUN /update_sources.sh
+
+RUN dpkg --add-architecture armhf
+RUN dpkg --add-architecture arm64
+RUN apt-get update && \
+    apt-get install -y \
+      debhelper \
+      dh-python \
+      python-all \
+      python-setuptools \
+      python-wheel \
+      python-numpy \
+      libpython-dev \
+      libpython-dev:armhf \
+      libpython-dev:arm64 \
+      python3-all \
+      python3-setuptools \
+      python3-wheel \
+      python3-numpy \
+      libpython3-dev \
+      libpython3-dev:armhf \
+      libpython3-dev:arm64 \
+      crossbuild-essential-armhf \
+      crossbuild-essential-arm64 \
+      zlib1g-dev  \
+      zlib1g-dev:armhf \
+      zlib1g-dev:arm64 \
+      swig \
+      curl \
+      unzip \
+      git && \
+    apt-get clean
diff --git a/tensorflow/lite/tools/pip_package/Dockerfile.debian b/tensorflow/lite/tools/pip_package/Dockerfile.debian
deleted file mode 100644
index 5a39e6bbe24..00000000000
--- a/tensorflow/lite/tools/pip_package/Dockerfile.debian
+++ /dev/null
@@ -1,31 +0,0 @@
-ARG VERSION
-FROM debian:${VERSION}
-
-RUN dpkg --add-architecture armhf
-RUN dpkg --add-architecture arm64
-RUN apt-get update && apt-get install -y \
-  debhelper \
-  dh-python \
-  python-all \
-  python-setuptools \
-  python-wheel \
-  python-numpy \
-  libpython-dev \
-  libpython-dev:armhf \
-  libpython-dev:arm64 \
-  python3-all \
-  python3-setuptools \
-  python3-wheel \
-  python3-numpy \
-  libpython3-dev \
-  libpython3-dev:armhf \
-  libpython3-dev:arm64 \
-  crossbuild-essential-armhf \
-  crossbuild-essential-arm64 \
-  zlib1g-dev  \
-  zlib1g-dev:armhf \
-  zlib1g-dev:arm64 \
-  swig \
-  curl \
-  unzip \
-  git
diff --git a/tensorflow/lite/tools/pip_package/Dockerfile.ubuntu b/tensorflow/lite/tools/pip_package/Dockerfile.ubuntu
deleted file mode 100644
index 35546af0691..00000000000
--- a/tensorflow/lite/tools/pip_package/Dockerfile.ubuntu
+++ /dev/null
@@ -1,34 +0,0 @@
-ARG VERSION
-FROM ubuntu:${VERSION}
-
-COPY update_sources.sh /
-RUN /update_sources.sh
-
-RUN dpkg --add-architecture armhf
-RUN dpkg --add-architecture arm64
-RUN apt-get update && apt-get install -y \
-  debhelper \
-  dh-python \
-  python-all \
-  python-setuptools \
-  python-wheel \
-  python-numpy \
-  libpython-dev \
-  libpython-dev:armhf \
-  libpython-dev:arm64 \
-  python3-all \
-  python3-setuptools \
-  python3-wheel \
-  python3-numpy \
-  libpython3-dev \
-  libpython3-dev:armhf \
-  libpython3-dev:arm64 \
-  crossbuild-essential-armhf \
-  crossbuild-essential-arm64 \
-  zlib1g-dev  \
-  zlib1g-dev:armhf \
-  zlib1g-dev:arm64 \
-  swig \
-  curl \
-  unzip \
-  git
diff --git a/tensorflow/lite/tools/pip_package/Makefile b/tensorflow/lite/tools/pip_package/Makefile
index eff94e274aa..97b7a627df4 100644
--- a/tensorflow/lite/tools/pip_package/Makefile
+++ b/tensorflow/lite/tools/pip_package/Makefile
@@ -9,9 +9,7 @@ BUILD_DEB ?= n
 
 MAKEFILE_DIR := $(realpath $(dir $(lastword $(MAKEFILE_LIST))))
 TENSORFLOW_DIR := $(MAKEFILE_DIR)/../../../..
-NAME := $(word 1, $(subst :, ,$(BASE_IMAGE)))
-VERSION := $(word 2, $(subst :, ,$(BASE_IMAGE)))
-TAG_IMAGE := "tflite-builder-$(BASE_IMAGE)"
+TAG_IMAGE := "tflite-runtime-builder-$(subst :,-,$(BASE_IMAGE))"
 
 .PHONY: help \
         docker-image \
@@ -26,10 +24,10 @@ help:
 	@echo "make clean        -- remove built wheel files"
 
 docker-image:
-	docker build -t $(TAG_IMAGE) --build-arg VERSION=$(VERSION) -f Dockerfile.$(NAME) .
+	docker build -t $(TAG_IMAGE) --build-arg IMAGE=$(BASE_IMAGE) .
 
 docker-shell: docker-image
-	docker run --rm -it -v $(TENSORFLOW_DIR):/tensorflow $(TAG_IMAGE)
+	docker run --rm -it -v $(TENSORFLOW_DIR):/tensorflow --workdir /tensorflow $(TAG_IMAGE)
 
 docker-build: docker-image
 	docker run \
@@ -40,7 +38,7 @@ docker-build: docker-image
 		-v $(CURDIR):/out \
 		--rm -it $(TAG_IMAGE) \
 		/bin/bash -c "tensorflow/tensorflow/lite/tools/pip_package/build_pip_package.sh && \
-		              cp /tmp/tflite_pip/*.deb /tmp/tflite_pip/$(PYTHON)/dist/*.whl /out"
+		              (cp /tmp/tflite_pip/*.deb /tmp/tflite_pip/$(PYTHON)/dist/*.whl /out 2>/dev/null || true)"
 
 clean:
-	rm -f *.whl
+	rm -f *.whl *.deb
diff --git a/tensorflow/lite/tools/pip_package/build_pip_package.sh b/tensorflow/lite/tools/pip_package/build_pip_package.sh
index 2a592190928..2e3e85f5fc2 100755
--- a/tensorflow/lite/tools/pip_package/build_pip_package.sh
+++ b/tensorflow/lite/tools/pip_package/build_pip_package.sh
@@ -68,5 +68,5 @@ if [[ "${TENSORFLOW_TARGET}" == "rpi" ]]; then
 elif [[ "${TENSORFLOW_TARGET}" == "aarch64" ]]; then
   dpkg-buildpackage -b -rfakeroot -us -uc -tc -d -a arm64
 else
-  dpkg-buildpackage -b -rfakeroot -us -uc -tc
+  dpkg-buildpackage -b -rfakeroot -us -uc -tc -d
 fi
diff --git a/tensorflow/lite/tools/pip_package/update_sources.sh b/tensorflow/lite/tools/pip_package/update_sources.sh
index 701a977e155..40e3213c597 100755
--- a/tensorflow/lite/tools/pip_package/update_sources.sh
+++ b/tensorflow/lite/tools/pip_package/update_sources.sh
@@ -14,10 +14,15 @@
 # limitations under the License.
 # ==============================================================================
 
+#!/bin/bash
 . /etc/os-release
 
+[[ "${NAME}" == "Ubuntu" ]] || exit 0
+
 sed -i "s/deb\ /deb \[arch=amd64\]\ /g" /etc/apt/sources.list
 
-echo "deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME} main universe" >> /etc/apt/sources.list
-echo "deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME}-updates main universe" >> /etc/apt/sources.list
-echo "deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME}-security main universe" >> /etc/apt/sources.list
+cat <<EOT >> /etc/apt/sources.list
+deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME} main universe
+deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME}-updates main universe
+deb [arch=arm64,armhf] http://ports.ubuntu.com/ubuntu-ports ${UBUNTU_CODENAME}-security main universe
+EOT