diff --git a/README.md b/README.md
index 9fdf808..83c85c5 100644
--- a/README.md
+++ b/README.md
@@ -197,7 +197,7 @@ The cache directories for different models correspond as follows:
 docker run -p 8080:8080 \
 -v /path/to/torch_cache:/root/.cache/torch \
 -v /path/to/huggingface_cache:/root/.cache/huggingface \
---rm lama-cleaner:cpu-0.21.0 \
+--rm lama-cleaner:cpu-0.24.4 \
 lama-cleaner --device=cpu --port=8080 --host=0.0.0.0
 ```
 
@@ -211,7 +211,7 @@ lama-cleaner --device=cpu --port=8080 --host=0.0.0.0
 docker run --gpus all -p 8080:8080 \
 -v /path/to/torch_cache:/root/.cache/torch \
 -v /path/to/huggingface_cache:/root/.cache/huggingface \
---rm lama-cleaner:cpu-0.21.0 \
+--rm lama-cleaner:gpu-0.24.4 \
 lama-cleaner --device=cuda --port=8080 --host=0.0.0.0
 ```
 
@@ -222,11 +222,11 @@ Then open [http://localhost:8080](http://localhost:8080)
 cpu only
 
 ```
-docker build -f ./docker/CPUDockerfile -t lamacleaner .
+docker build -f --build-arg version=0.x.0 ./docker/CPUDockerfile -t lamacleaner .
 ```
 
 gpu & cpu
 
 ```
-docker build -f ./docker/GPUDockerfile -t lamacleaner .
+docker build -f --build-arg version=0.x.0 ./docker/GPUDockerfile -t lamacleaner .
 ```
diff --git a/build_docker.sh b/build_docker.sh
index 9a68de2..ac0d2b9 100644
--- a/build_docker.sh
+++ b/build_docker.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -e
 
-GIT_TAG=`git describe --tags --abbrev=0`
+GIT_TAG=$1
 IMAGE_DESC="Image inpainting tool powered by SOTA AI Model" 
 GIT_REPO="https://github.com/Sanster/lama-cleaner"
 
@@ -14,6 +14,7 @@ docker buildx build \
 --label org.opencontainers.image.url=$GIT_REPO \
 --label org.opencontainers.image.source=$GIT_REPO \
 --label org.opencontainers.image.version=$GIT_TAG \
+--build-arg version=$GIT_TAG \
 --tag lama-cleaner:cpu-$GIT_TAG .
 
 
@@ -26,4 +27,5 @@ docker buildx build \
 --label org.opencontainers.image.url=$GIT_REPO \
 --label org.opencontainers.image.source=$GIT_REPO \
 --label org.opencontainers.image.version=$GIT_TAG \
+--build-arg version=$GIT_TAG \
 --tag lama-cleaner:gpu-$GIT_TAG .
diff --git a/docker/CPUDockerfile b/docker/CPUDockerfile
index bd7bfbe..33bb7bb 100644
--- a/docker/CPUDockerfile
+++ b/docker/CPUDockerfile
@@ -7,8 +7,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 RUN pip install --upgrade pip && \
     pip install torch==1.12.1 --extra-index-url https://download.pytorch.org/whl/cpu
-RUN pip install lama-cleaner
+
+ARG version
+
+RUN pip install lama-cleaner==$version
 
 EXPOSE 8080
 
-CMD ["bash"]
\ No newline at end of file
+CMD ["bash"]
diff --git a/docker/GPUDockerfile b/docker/GPUDockerfile
index 8680366..f6b9264 100644
--- a/docker/GPUDockerfile
+++ b/docker/GPUDockerfile
@@ -6,9 +6,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     curl python3-pip
 
 RUN pip3 install --upgrade pip && \
-    pip3 install torch==1.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
-RUN pip3 install lama-cleaner
+    pip3 install torch==1.12.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+    
+ARG version 
+
+RUN pip3 install lama-cleaner==$version
 
 EXPOSE 8080
 
-CMD ["bash"]
\ No newline at end of file
+CMD ["bash"]