diff --git a/setup.sh b/setup.sh
index 21ae933..479faad 100644
--- a/setup.sh
+++ b/setup.sh
@@ -194,10 +194,14 @@ function select_part () {
     part_type="${1:?}" # 'boot' or 'zfs'
     zfs_install_drive="${2:-}"
 
-    if [[ "${part_schema}" = 'mbr' ]]; then
-        specific_part_type='boot'
+    if [[ "${part_type}" = 'boot' ]]; then
+        if [[ "${part_schema}" = 'mbr' ]]; then
+            specific_part_type='boot'
+        else
+            specific_part_type='efi'
+        fi
     else
-        specific_part_type='efi'
+        specific_part_type="${part_type}"
     fi
 
     if [[ "${zfs_install_drive}" ]]; then