diff --git a/src/factory.ts b/src/factory.ts
index 5899b36..e8218d6 100644
--- a/src/factory.ts
+++ b/src/factory.ts
@@ -284,7 +284,7 @@ export async function flashZip(
             "reboot",
             "device",
             FASTBOOTD_REBOOT_TIME,
-            device.reboot("fastboot", true, onReconnect)
+            tryReboot(device, "fastboot", onReconnect)
         );
 
         let superName = await device.getVariable("super-partition-name");
@@ -322,7 +322,7 @@ export async function flashZip(
             "reboot",
             "device",
             BOOTLOADER_REBOOT_TIME,
-            device.reboot("bootloader", true, onReconnect)
+            tryReboot(device, "bootloader", onReconnect)
         );
     }