diff --git a/Dockerfile b/Dockerfile
index 03a4c307..c35de629 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -39,7 +39,7 @@ RUN pip install --upgrade pip
 ################################################################################
 
 # Leverage multi-stage build for unit tests.
-FROM base as ci
+FROM base AS ci
 
 # TODO(markblee): Remove gcp,vertexai_tensorboard from CI.
 RUN pip install .[core,dev,grain,gcp,vertexai_tensorboard]