diff --git a/comps/llms/text-generation/vllm/docker/Dockerfile.microservice b/comps/llms/text-generation/vllm/docker/Dockerfile.microservice
index 765b8b33e5..af5a0cc196 100644
--- a/comps/llms/text-generation/vllm/docker/Dockerfile.microservice
+++ b/comps/llms/text-generation/vllm/docker/Dockerfile.microservice
@@ -1,9 +1,10 @@
-
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 FROM langchain/langchain:latest
 
+ARG ARCH="cpu"  # Set this to "cpu" or "gpu"
+
 RUN apt-get update -y && apt-get install -y --no-install-recommends --fix-missing \
     libgl1-mesa-glx \
     libjemalloc-dev \
@@ -18,7 +19,11 @@ USER user
 COPY comps /home/user/comps
 
 RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir -r /home/user/comps/llms/text-generation/vllm/requirements.txt
+    if [ ${ARCH} = "cpu" ]; then \
+      pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu -r /home/user/comps/llms/text-generation/vllm/requirements.txt; \
+    else \
+      pip install --no-cache-dir -r /home/user/comps/llms/text-generation/vllm/requirements.txt; \
+    fi
 
 ENV PYTHONPATH=$PYTHONPATH:/home/user