diff --git a/Dockerfile.ubi b/Dockerfile.ubi
index 294399be24c46..27271c83566e6 100644
--- a/Dockerfile.ubi
+++ b/Dockerfile.ubi
@@ -169,7 +169,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # We used base cuda image because pytorch installs its own cuda libraries.
 # However pynccl depends on cuda libraries so we had to switch to the runtime image
 # In the future it would be nice to get a container with pytorch and cuda without duplicating cuda
-FROM python-install AS vllm-openai
+FROM cuda-base AS vllm-openai
 
 WORKDIR /workspace