diff --git a/pmb/config/__init__.py b/pmb/config/__init__.py
index a22640c457e5f3733926e3a1d09809de3dcf05d1..2313d2cefab4577c4ab8f9915d9ce03da3460c0b 100644
--- a/pmb/config/__init__.py
+++ b/pmb/config/__init__.py
@@ -419,8 +419,7 @@ flashers: dict[str, dict[str, bool | list[str] | dict[str, list[list[str]]]]] =
         "actions": {
             "list_devices": [["heimdall", "detect"]],
             "flash_rootfs": [
-                ["heimdall_wait_for_device.sh"],
-                ["heimdall", "flash", "--$PARTITION_ROOTFS", "$IMAGE"],
+                ["heimdall", "flash", "--wait", "--$PARTITION_ROOTFS", "$IMAGE"],
             ],
             "flash_kernel": [
                 [
@@ -440,14 +439,21 @@ flashers: dict[str, dict[str, bool | list[str] | dict[str, list[list[str]]]]] =
         "actions": {
             "list_devices": [["heimdall", "detect"]],
             "flash_rootfs": [
-                ["heimdall_wait_for_device.sh"],
-                ["heimdall", "flash", "--$PARTITION_ROOTFS", "$IMAGE", "$NO_REBOOT", "$RESUME"],
+                [
+                    "heimdall",
+                    "flash",
+                    "--wait",
+                    "--$PARTITION_ROOTFS",
+                    "$IMAGE",
+                    "$NO_REBOOT",
+                    "$RESUME",
+                ],
             ],
             "flash_kernel": [
-                ["heimdall_wait_for_device.sh"],
                 [
                     "heimdall",
                     "flash",
+                    "--wait",
                     "--$PARTITION_KERNEL",
                     "$BOOT/boot.img$FLAVOR",
                     "$NO_REBOOT",
@@ -476,10 +482,10 @@ flashers: dict[str, dict[str, bool | list[str] | dict[str, list[list[str]]]]] =
                 ["rm", "-f", "/vbmeta.img"],
             ],
             "flash_lk2nd": [
-                ["heimdall_wait_for_device.sh"],
                 [
                     "heimdall",
                     "flash",
+                    "--wait",
                     "--$PARTITION_KERNEL",
                     "$BOOT/lk2nd.img",
                     "$NO_REBOOT",