diff --git a/images/bloomz-560m/Dockerfile b/images/bloomz-560m/Dockerfile
index f715ece..e7e2e8d 100644
--- a/images/bloomz-560m/Dockerfile
+++ b/images/bloomz-560m/Dockerfile
@@ -53,9 +53,10 @@ RUN update-alternatives --install /usr/bin/python python ${ENVD_PREFIX}/python 1
     update-alternatives --install /usr/bin/pip pip ${ENVD_PREFIX}/pip 1 && \
     update-alternatives --install /usr/bin/pip3 pip3 ${ENVD_PREFIX}/pip3 1
 
-COPY requirements.txt /
+COPY requirements-cpu.txt /
 
-RUN pip install -r requirements.txt
+RUN pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+RUN pip install -r requirements-cpu.txt
 
 RUN mkdir -p /workspace
 
diff --git a/requirements-cpu.txt b/requirements-cpu.txt
index f217def..2078e2a 100644
--- a/requirements-cpu.txt
+++ b/requirements-cpu.txt
@@ -1,6 +1,6 @@
 msgpack
 mosec
-torch --extra-index-url https://download.pytorch.org/whl/cpu
+torch
 transformers
 llmspec
 falcon