diff --git a/samples/livepatch/livepatch-sample.c b/samples/livepatch/livepatch-sample.c
index 2ffeae67dc1454dbc20b042c5ec0e426f64b5da0..915f717586b785a598006e02b8057970ba91ab23 100644
--- a/samples/livepatch/livepatch-sample.c
+++ b/samples/livepatch/livepatch-sample.c
@@ -33,6 +33,31 @@
  */
 
 #include <linux/seq_file.h>
+
+#ifdef CONFIG_LIVEPATCH_STOP_MACHINE_CONSISTENCY
+void load_hook(void)
+{
+	pr_info("load_hook\n");
+}
+
+void unload_hook(void)
+{
+	pr_info("unload_hook\n");
+}
+
+static struct klp_hook hooks_load[] = {
+	{
+		.hook = load_hook
+	}, { }
+};
+
+static struct klp_hook hooks_unload[] = {
+	{
+		.hook = unload_hook
+	}, { }
+};
+#endif
+
 static int livepatch_cmdline_proc_show(struct seq_file *m, void *v)
 {
 	seq_printf(m, "%s\n", "this has been live patched");
@@ -54,6 +79,10 @@ static struct klp_object objs[] = {
 	{
 		/* name being NULL means vmlinux */
 		.funcs = funcs,
+#ifdef CONFIG_LIVEPATCH_STOP_MACHINE_CONSISTENCY
+		.hooks_load = hooks_load,
+		.hooks_unload = hooks_unload,
+#endif
 	}, { }
 };
 
@@ -68,11 +97,19 @@ static int livepatch_init(void)
 	patch.objs[0].funcs[0].new_func =
 		(void *)ppc_function_entry((void *)livepatch_cmdline_proc_show);
 #endif
+
+#ifdef CONFIG_LIVEPATCH_PER_TASK_CONSISTENCY
 	return klp_enable_patch(&patch);
+#elif defined(CONFIG_LIVEPATCH_STOP_MACHINE_CONSISTENCY)
+	return klp_register_patch(&patch);
+#endif
 }
 
 static void livepatch_exit(void)
 {
+#ifdef CONFIG_LIVEPATCH_STOP_MACHINE_CONSISTENCY
+	WARN_ON(klp_unregister_patch(&patch));
+#endif
 }
 
 module_init(livepatch_init);