diff --git a/build_library/base_image_util.sh b/build_library/base_image_util.sh
index 56ff4eef61..9006021be6 100755
--- a/build_library/base_image_util.sh
+++ b/build_library/base_image_util.sh
@@ -39,11 +39,9 @@ cleanup_mounts() {
 create_base_image() {
   local image_name=$1
   local rootfs_verification_enabled=$2
-  local image_type="base"
 
-  if [[ "${FLAGS_disk_layout}" != "default" ]]; then
-      image_type="${FLAGS_disk_layout}"
-  fi
+  get_disk_layout_type
+  local image_type="${DISK_LAYOUT_TYPE}"
 
   check_valid_layout "base"
   check_valid_layout ${image_type}
diff --git a/build_library/disk_layout_util.sh b/build_library/disk_layout_util.sh
index 56d8cb2a4c..e337111e70 100644
--- a/build_library/disk_layout_util.sh
+++ b/build_library/disk_layout_util.sh
@@ -112,6 +112,10 @@ check_valid_layout() {
 
 get_disk_layout_type() {
   DISK_LAYOUT_TYPE="base"
+  if [[ -n "${FLAGS_disk_layout}" && \
+      "${FLAGS_disk_layout}" != "default" ]]; then
+    DISK_LAYOUT_TYPE="${FLAGS_disk_layout}"
+  fi
 }
 
 emit_gpt_scripts() {