diff --git a/build.py b/build.py
index 716b227c42..38afd4fd05 100755
--- a/build.py
+++ b/build.py
@@ -1239,10 +1239,10 @@ def create_dockerfile_linux(
 
 WORKDIR /opt/tritonserver
 COPY --chown=1000:1000 NVIDIA_Deep_Learning_Container_License.pdf .
-RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \\
-    "tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[all] && \\
-    find /opt/tritonserver/python -maxdepth 1 -type f -name \\
-    "tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[all]
+RUN find /opt/tritonserver/python -maxdepth 1 -type f -name \
+    "tritonserver-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT] && \
+    find /opt/tritonserver/python -maxdepth 1 -type f -name \
+    "tritonfrontend-*.whl" | xargs -I {} pip install --upgrade {}[$VARIANT];
 
 RUN pip3 install -r python/openai/requirements.txt
 
@@ -1295,6 +1295,7 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
     df = """
 ARG TRITON_VERSION
 ARG TRITON_CONTAINER_VERSION
+ARG VARIANT=all
 
 ENV TRITON_SERVER_VERSION ${TRITON_VERSION}
 ENV NVIDIA_TRITON_SERVER_VERSION ${TRITON_CONTAINER_VERSION}
@@ -1898,6 +1899,10 @@ def create_docker_build_script(script_name, container_install_dir, container_ci_
                 f"--secret id=NVPL_SLIM_URL",
                 f"--build-arg BUILD_PUBLIC_VLLM={build_public_vllm}",
             ]
+        if  FLAGS.build_variant:
+            finalargs += [
+                "--build-arg VARIANT="+(FLAGS.build_variant),
+            ]
         finalargs += [
             "-t",
             "tritonserver",
@@ -2763,6 +2768,13 @@ def enable_all():
         default=DEFAULT_TRITON_VERSION_MAP["rhel_py_version"],
         help="This flag sets the Python version for RHEL platform of Triton Inference Server to be built. Default: the latest supported version.",
     )
+    parser.add_argument(
+        "--build_variant",
+        required=False,
+        type=str,
+        default="all",
+        help="Can be set to all or cpu,Default value is all."
+    )
     parser.add_argument(
         "--build-secret",
         action="append",
@@ -2802,6 +2814,9 @@ def enable_all():
         FLAGS.extra_backend_cmake_arg = []
     if FLAGS.build_secret is None:
         FLAGS.build_secret = []
+    if hasattr(FLAGS, 'build_variant') and FLAGS.build_variant not in ["all", "cpu"]:
+        raise ValueError(f"Invalid build_variant value: {FLAGS.build_variant}. Expected 'all' or 'cpu'.")
+
 
     # if --enable-all is specified, then update FLAGS to enable all
     # settings, backends, repo-agents, caches, file systems, endpoints, etc.