diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 4223002d80b4..4bfe63a5225d 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -58,13 +58,6 @@ static const struct block_device_operations zram_devops;
 
 static void slot_free(struct zram *zram, u32 index);
 
-static void slot_lock_init(struct zram *zram)
-{
-	static struct lock_class_key __key;
-
-	lockdep_init_map(&zram->table_lock_map, "zram->table[index].lock", &__key, 0);
-}
-
 /*
  * entry locking rules:
  *
@@ -1978,6 +1971,7 @@ static void zram_meta_free(struct zram *zram, u64 disksize)
 	zs_destroy_pool(zram->mem_pool);
 	vfree(zram->table);
 	zram->table = NULL;
+	lockdep_unregister_key(&zram->table_lock_key);
 }
 
 static bool zram_meta_alloc(struct zram *zram, u64 disksize)
@@ -1999,7 +1993,8 @@ static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 	if (!huge_class_size)
 		huge_class_size = zs_huge_class_size(zram->mem_pool);
 
-	slot_lock_init(zram);
+	lockdep_register_key(&zram->table_lock_key);
+	lockdep_init_map(&zram->table_lock_map, "zram->table[index].lock", &zram->table_lock_key, 0);
 
 	return true;
 }
diff --git a/drivers/block/zram/zram_drv.h b/drivers/block/zram/zram_drv.h
index 8290dc6e2335..4fddc582f3b8 100644
--- a/drivers/block/zram/zram_drv.h
+++ b/drivers/block/zram/zram_drv.h
@@ -107,6 +107,7 @@ struct zram_stats {
 struct zram {
 	struct zram_table_entry *table;
 	struct lockdep_map table_lock_map;
+	struct lock_class_key table_lock_key;
 	struct zs_pool *mem_pool;
 	struct zcomp *comps[ZRAM_MAX_COMPS];
 	struct zcomp_params params[ZRAM_MAX_COMPS];