diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu
index 8bc32e30..6156bd95 100644
--- a/docker/Dockerfile.ubuntu
+++ b/docker/Dockerfile.ubuntu
@@ -24,8 +24,6 @@ RUN pip --no-cache-dir install omegaconf==2.0.6
 
 RUN pip install --no-cache-dir 'cython<=0.29.36' kwsbp==0.0.6 versioneer numpy==1.25.1 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
 
-RUN echo "cache bust $(date +%Y%m%d%H%M%S)"
-
 # install dependencies
 COPY requirements /var/modelscope
 RUN pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
@@ -40,10 +38,14 @@ RUN pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://model
     pip install --no-cache-dir 'scipy<1.13.0' && \
     pip cache purge
 
+RUN echo "cache bust $(date +%Y%m%d%H%M%S)"
+
 RUN cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b  {modelscope_branch}  --single-branch https://github.com/modelscope/modelscope.git && cd modelscope && pip install . -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && cd / && rm -fr /tmp/modelscope && pip cache purge;
 
 RUN cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b {swift_branch}  --single-branch https://github.com/modelscope/ms-swift.git && cd ms-swift && pip install .[all] && cd / && rm -fr /tmp/ms-swift && pip cache purge;
 
+pip install --no-cache-dir torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version} {index_url}
+
 RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
     pip config set install.trusted-host mirrors.aliyun.com && \
     cp /tmp/resources/ubuntu2204.aliyun /etc/apt/sources.list
diff --git a/docker/build_image.py b/docker/build_image.py
index 76b50688..ed8e9b94 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -147,8 +147,10 @@ class CPUImageBuilder(Builder):
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{meta_file}', meta_file)
             content = content.replace('{version_args}', version_args)
-            content = content.replace('{resolver}',
-                                      '--use-deprecated=legacy-resolver')
+            content = content.replace('{torch_version}', self.args.torch_version)
+            content = content.replace('{torchvision_version}', self.args.torchvision_version)
+            content = content.replace('{torchaudio_version}', self.args.torchaudio_version)
+            content = content.replace('{index_url}', '--index-url https://download.pytorch.org/whl/cpu')
             content = content.replace('{modelscope_branch}',
                                       self.args.modelscope_branch)
             content = content.replace('{swift_branch}', self.args.swift_branch)
@@ -197,7 +199,10 @@ class GPUImageBuilder(Builder):
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{meta_file}', meta_file)
             content = content.replace('{version_args}', version_args)
-            content = content.replace('{resolver}', '')
+            content = content.replace('{torch_version}', self.args.torch_version)
+            content = content.replace('{torchvision_version}', self.args.torchvision_version)
+            content = content.replace('{torchaudio_version}', self.args.torchaudio_version)
+            content = content.replace('{index_url}', '')
             content = content.replace('{modelscope_branch}',
                                       self.args.modelscope_branch)
             content = content.replace('{swift_branch}', self.args.swift_branch)
@@ -264,7 +269,10 @@ class LLMImageBuilder(Builder):
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{meta_file}', meta_file)
             content = content.replace('{version_args}', version_args)
-            content = content.replace('{resolver}', '')
+            content = content.replace('{torch_version}', self.args.torch_version)
+            content = content.replace('{torchvision_version}', self.args.torchvision_version)
+            content = content.replace('{torchaudio_version}', self.args.torchaudio_version)
+            content = content.replace('{index_url}', '')
             content = content.replace('{modelscope_branch}',
                                       self.args.modelscope_branch)
             content = content.replace('{swift_branch}', self.args.swift_branch)
diff --git a/docker/install.sh b/docker/install.sh
index 145b43ad..0daecf2e 100644
--- a/docker/install.sh
+++ b/docker/install.sh
@@ -46,5 +46,3 @@ pip install --no-cache-dir vllm==$vllm_version -U
 pip install --no-cache-dir -U lmdeploy==$lmdeploy_version --no-deps
 
 pip install --no-cache-dir pynvml shortuuid
-
-pip install --no-cache-dir torch==$torch_version torchvision==$torchvision_version torchaudio==$torchaudio_version
diff --git a/docker/install_cpu.sh b/docker/install_cpu.sh
index 1cd3396b..fce62afe 100644
--- a/docker/install_cpu.sh
+++ b/docker/install_cpu.sh
@@ -19,5 +19,3 @@ pip install --no-cache-dir -U qwen_vl_utils pyav librosa timm transformers accel
 pip install --no-cache-dir text2sql_lgesql==1.3.0 git+https://github.com/jin-s13/xtcocoapi.git@v1.14 git+https://github.com/gatagat/lap.git@v0.4.0 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force --no-deps
 
 pip install --no-cache-dir mpi4py paint_ldm mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 ipykernel fasttext -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
-
-pip install --no-cache-dir -U torch==$torch_version torchvision==$torchvision_version torchaudio==$torchaudio_version --index-url https://download.pytorch.org/whl/cpu