diff --git a/setup.sh b/setup.sh
index 193882c..dfc80c2 100644
--- a/setup.sh
+++ b/setup.sh
@@ -61,11 +61,7 @@ function get_parts () {
             parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.children[].parttype=="6a85cf4d-1dd2-11b2-99a6-080020736631") | .children[].path')" || exit 1
             ;;
         efi)
-            if [[ "${zfs_install_drive}" ]]; then
-                parts="$(get_partitions | jq --raw-output '.[][] | select(.path=="'"${zfs_install_drive}"'") | select(.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .path')" || exit 1
-            else
-                parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.children[].parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .children[].path')" || exit 1
-            fi
+            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.children[].path=="'"${zfs_install_drive:?}"'") | select(.children[].parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .children[].path')" || exit 1
             ;;
         *)
             >2 printf -- '%s\n' 'Unknown partition type '"'"'"${parttype}"'"'"', exiting ...'
@@ -117,6 +113,7 @@ function select_part () {
     part_type="${1:?}" # 'efi' or 'zfs'
     zfs_install_drive="${2:-}"
     if [[ "${zfs_install_drive}" ]]; then
+        # This is intended to find correct EFI partition
         parts="$(get_parts "${part_type}" "${zfs_install_drive}")"
     else
         parts="$(get_parts "${part_type}")"