diff --git a/docker/Dockerfile b/docker/Dockerfile
index 97399c939..a1e136ee5 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -21,7 +21,10 @@ RUN apt-get update && \
 RUN conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
 
 # install ninja
-RUN apt-get install -y --no-install-recommends ninja-build
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends ninja-build && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # install apex
 RUN git clone https://github.com/NVIDIA/apex && \
@@ -31,7 +34,7 @@ RUN git clone https://github.com/NVIDIA/apex && \
     pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" ./
 
 # install colossalai
-ARG VERSION=1
+ARG VERSION=main
 RUN git clone -b ${VERSION} https://github.com/hpcaitech/ColossalAI.git \
     && cd ./ColossalAI \
     && CUDA_EXT=1 pip install -v --no-cache-dir .