diff --git a/Sparserestore/__init__.py b/Sparserestore/__init__.py
index 0c037de..2a94512 100644
--- a/Sparserestore/__init__.py
+++ b/Sparserestore/__init__.py
@@ -9,6 +9,13 @@ from pymobiledevice3.lockdown import LockdownClient
 
 from . import backup
 
+def reboot_device(reboot: bool = False, lockdown_client: LockdownClient = None):
+    if reboot and lockdown_client != None:
+        print("Success! Rebooting your device...")
+        with DiagnosticsService(lockdown_client) as diagnostics_service:
+            diagnostics_service.restart()
+        print("Remember to turn Find My back on!")
+
 def perform_restore(backup: backup.Backup, reboot: bool = False, lockdown_client: LockdownClient = None):
     try:
         with TemporaryDirectory() as backup_dir:
@@ -18,6 +25,8 @@ def perform_restore(backup: backup.Backup, reboot: bool = False, lockdown_client
                 lockdown_client = create_using_usbmux()
             with Mobilebackup2Service(lockdown_client) as mb:
                 mb.restore(backup_dir, system=True, reboot=False, copy=False, source=".")
+            # reboot the device
+            reboot_device(reboot, lockdown_client)
     except PyMobileDevice3Exception as e:
         if "Find My" in str(e):
             print("Find My must be disabled in order to use this tool.")
@@ -26,8 +35,4 @@ def perform_restore(backup: backup.Backup, reboot: bool = False, lockdown_client
         elif "crash_on_purpose" not in str(e):
             raise e
         else:
-            if reboot and lockdown_client != None:
-                print("Success! Rebooting your device...")
-                with DiagnosticsService(lockdown_client) as diagnostics_service:
-                    diagnostics_service.restart()
-                print("Remember to turn Find My back on!")
\ No newline at end of file
+            reboot_device(reboot, lockdown_client)
\ No newline at end of file