From 9774f513e607d4c5d0236bf62be34b90be041c2c Mon Sep 17 00:00:00 2001
From: "yuze.zyz" <yuze.zyz@alibaba-inc.com>
Date: Fri, 25 Oct 2024 12:48:23 +0800
Subject: [PATCH] fix

---
 docker/Dockerfile.ubuntu |  2 +-
 docker/build_image.py    | 31 +++++++++++++++++++++----------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu
index 6156bd95..056c0026 100644
--- a/docker/Dockerfile.ubuntu
+++ b/docker/Dockerfile.ubuntu
@@ -44,7 +44,7 @@ RUN cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b  {modelscope_branch}  --single
 
 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 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 && \
diff --git a/docker/build_image.py b/docker/build_image.py
index ed8e9b94..d53708eb 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -147,10 +147,15 @@ 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('{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('{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)
@@ -199,9 +204,12 @@ 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('{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('{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)
@@ -269,9 +277,12 @@ 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('{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('{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)