diff --git a/docker/build.sh b/docker/build.sh
index 4d3fdd85e6..c5670aa29b 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -35,7 +35,7 @@ else
 fi
 
 # Build Container
-docker build \
+DOCKER_BUILDKIT=1 docker build \
     --platform="${PLATFORM}" \
     --tag="${CONTAINER_IMAGE}" \
     ${PIP_EXTRA_INDEX_URL:+--build-arg="PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}"} \
diff --git a/docker/env.sh b/docker/env.sh
index eab878577e..2196d78425 100644
--- a/docker/env.sh
+++ b/docker/env.sh
@@ -7,7 +7,7 @@ if [[ -z "$PIP_EXTRA_INDEX_URL" ]]; then
     CUDA_AVAILABLE=$(python -c "import torch;print(torch.cuda.is_available())")
     ROCM_AVAILABLE=$(python -c "import torch;print(torch.version.hip is not None)")
     if  [[ "$(uname -s)" != "Darwin" && "${CUDA_AVAILABLE}" == "True" ]]; then
-      CONTAINER_FLAVOR=cuda
+      CONTAINER_FLAVOR="cuda"
     elif [[ "$(uname -s)" != "Darwin" && "${ROCM_AVAILABLE}" == "True" ]]; then
       CONTAINER_FLAVOR="rocm"
     else
@@ -16,9 +16,9 @@ if [[ -z "$PIP_EXTRA_INDEX_URL" ]]; then
   fi
   # Set PIP_EXTRA_INDEX_URL based on container flavor
   if [[ "$CONTAINER_FLAVOR" == "rocm" ]]; then
-    PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL-"https://download.pytorch.org/whl/rocm"}"
-  elif CONTAINER_FLAVOR=cpu; then
-    PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL-"https://download.pytorch.org/whl/cpu"}"
+    PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/rocm"
+  elif [[ "$CONTAINER_FLAVOR" == "cpu" ]]; then
+    PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
   fi
 fi