{
        cpu_pm_register_notifier(&hyp_init_cpu_pm_nb);
 }
+static void __init hyp_cpu_pm_exit(void)
+{
+       cpu_pm_unregister_notifier(&hyp_init_cpu_pm_nb);
+}
 #else
 static inline void hyp_cpu_pm_init(void)
 {
 }
+static inline void hyp_cpu_pm_exit(void)
+{
+}
 #endif
 
 static void teardown_common_resources(void)
        /*
         * Register CPU Hotplug notifier
         */
-       cpu_notifier_register_begin();
-       err = __register_cpu_notifier(&hyp_init_cpu_nb);
-       cpu_notifier_register_done();
+       err = register_cpu_notifier(&hyp_init_cpu_nb);
        if (err) {
                kvm_err("Cannot register KVM init CPU notifier (%d)\n", err);
                return err;
        free_hyp_pgds();
        for_each_possible_cpu(cpu)
                free_page(per_cpu(kvm_arm_hyp_stack_page, cpu));
+       unregister_cpu_notifier(&hyp_init_cpu_nb);
+       hyp_cpu_pm_exit();
 }
 
 static int init_vhe_mode(void)