.fops   =       &wdt_fops,
 };
 
+static int wdt_restart_handle(struct notifier_block *this, unsigned long mode,
+                             void *cmd)
+{
+       /*
+        * Cobalt devices have no way of rebooting themselves other
+        * than getting the watchdog to pull reset, so we restart the
+        * watchdog on reboot with no heartbeat.
+        */
+       wdt_change(WDT_ENABLE);
+
+       /* loop until the watchdog fires */
+       while (true)
+               ;
+
+       return NOTIFY_DONE;
+}
+
+static struct notifier_block wdt_restart_handler = {
+       .notifier_call = wdt_restart_handle,
+       .priority = 128,
+};
+
 /*
  *     Notifier for system down
  */
        if (code == SYS_DOWN || code == SYS_HALT)
                wdt_turnoff();
 
-       if (code == SYS_RESTART) {
-               /*
-                * Cobalt devices have no way of rebooting themselves other
-                * than getting the watchdog to pull reset, so we restart the
-                * watchdog on reboot with no heartbeat
-                */
-               wdt_change(WDT_ENABLE);
-               pr_info("Watchdog timer is now enabled with no heartbeat - should reboot in ~1 second\n");
-       }
        return NOTIFY_DONE;
 }
 
        /* Deregister */
        misc_deregister(&wdt_miscdev);
        unregister_reboot_notifier(&wdt_notifier);
+       unregister_restart_handler(&wdt_restart_handler);
        pci_dev_put(alim7101_pmu);
 }
 
                goto err_out;
        }
 
+       rc = register_restart_handler(&wdt_restart_handler);
+       if (rc) {
+               pr_err("cannot register restart handler (err=%d)\n", rc);
+               goto err_out_reboot;
+       }
+
        rc = misc_register(&wdt_miscdev);
        if (rc) {
                pr_err("cannot register miscdev on minor=%d (err=%d)\n",
                       wdt_miscdev.minor, rc);
-               goto err_out_reboot;
+               goto err_out_restart;
        }
 
        if (nowayout)
                timeout, nowayout);
        return 0;
 
+err_out_restart:
+       unregister_restart_handler(&wdt_restart_handler);
 err_out_reboot:
        unregister_reboot_notifier(&wdt_notifier);
 err_out: