1 files changed, 4 insertions, 4 deletions

diff --git a/mm/zswap.c b/mm/zswap.c
index e55bab9dc41f..d7337fbf6605 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -387,18 +387,18 @@ static int zswap_cpu_init(void)
 {
         unsigned long cpu;
 
-        get_online_cpus();
+        cpu_notifier_register_begin();
         for_each_online_cpu(cpu)
                 if (__zswap_cpu_notifier(CPU_UP_PREPARE, cpu) != NOTIFY_OK)
                         goto cleanup;
-        register_cpu_notifier(&zswap_cpu_notifier_block);
-        put_online_cpus();
+        __register_cpu_notifier(&zswap_cpu_notifier_block);
+        cpu_notifier_register_done();
         return 0;
 
 cleanup:
         for_each_online_cpu(cpu)
                 __zswap_cpu_notifier(CPU_UP_CANCELED, cpu);
-        put_online_cpus();
+        cpu_notifier_register_done();
         return -ENOMEM;
 }