diff --git a/Dockerfile b/Dockerfile
index 0e8089802..b9f88c029 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -52,12 +52,19 @@ RUN sudo apt-get update && \
     sudo chown -R "$(whoami):$(whoami)" /usr/lib/go-1.16 && \
     sudo rm -rf /var/lib/apt/lists/*
 
+# Install addlicense
 ENV GOPATH="/usr/lib/go-1.16"
 ENV GOBIN="${GOPATH}/bin"
 ENV GOROOT="${GOPATH}"
 ENV PATH="${GOBIN}:${PATH}"
 RUN go install github.com/google/addlicense@latest
 
+# Install extra PyPI dependencies
+COPY --chown=torchopt tests/requirements.txt tests/requirements.txt
+RUN source ~/venv/bin/activate && \
+    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" -r tests/requirements.txt && \
+    rm -rf ~/.pip/cache ~/.cache/pip
+
 ####################################################################################################
 
 FROM builder AS base
@@ -75,16 +82,4 @@ ENTRYPOINT [ "/bin/bash", "--login" ]
 
 FROM devel-builder AS devel
 
-COPY --chown=torchopt . .
-
-# Install extra dependencies
-RUN source ~/venv/bin/activate && \
-    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" -r tests/requirements.txt && \
-    rm -rf ~/.pip/cache ~/.cache/pip
-
-# Install TorchOpt
-RUN source ~/venv/bin/activate && \
-    python -m pip install -e . && \
-    rm -rf .eggs *.egg-info ~/.pip/cache ~/.cache/pip
-
-ENTRYPOINT [ "/bin/bash", "--login" ]
+COPY --from=base /home/torchopt/TorchOpt .
diff --git a/Makefile b/Makefile
index b4b1e2628..775ff7e78 100644
--- a/Makefile
+++ b/Makefile
@@ -136,23 +136,15 @@ clean: clean-py clean-build clean-docs
 
 # Build docker images
 
-docker-dev:
-	docker build --network=host -t $(PROJECT_NAME):$(COMMIT_HASH) -f docker/dev.dockerfile .
-	docker run --network=host -v /:/host -it $(PROJECT_NAME):$(COMMIT_HASH) bash
+docker-base:
+	docker build --target base --tag $(PROJECT_NAME):$(COMMIT_HASH) --file Dockerfile .
 	echo successfully build docker image with tag $(PROJECT_NAME):$(COMMIT_HASH)
 
-docker-release:
-	docker build --network=host -t $(PROJECT_NAME)-release:$(COMMIT_HASH) -f docker/release.dockerfile .
-	mkdir -p wheelhouse
-	docker run --network=host -v `pwd`/wheelhouse:/whl -it $(PROJECT_NAME)-release:$(COMMIT_HASH) bash -c "cp wheelhouse/* /whl"
-	echo successfully build docker image with tag $(PROJECT_NAME)-release:$(COMMIT_HASH)
+docker-devel:
+	docker build --target devel --tag $(PROJECT_NAME)-devel:$(COMMIT_HASH) --file Dockerfile .
+	echo successfully build docker image with tag $(PROJECT_NAME)-devel:$(COMMIT_HASH)
 
+docker: docker-base docker-devel
 
-docker-test:
-	sudo docker build --network=host -t $(PROJECT_NAME):$(COMMIT_HASH) -f docker/test.dockerfile .
-	sudo docker run --network=host -v /:/host -it $(PROJECT_NAME):$(COMMIT_HASH) bash
-	echo successfully build docker image with tag $(PROJECT_NAME):$(COMMIT_HASH)
-
-
-pypi-wheel: 
-	python -m pip install --upgrade pip && pip install setuptools wheel twine && python setup.py sdist bdist_wheel &&twine upload dist/*
+docker-run-devel: docker-devel
+	docker run --network=host --gpus=all -v /:/host -h ubuntu -it $(PROJECT_NAME)-devel:$(COMMIT_HASH)