diff --git a/src/tree_data_common.c b/src/tree_data_common.c
index 65784ae..4602e98 100644
--- a/src/tree_data_common.c
+++ b/src/tree_data_common.c
@@ -748,8 +748,8 @@
     assert(!node->schema);
 
     /* get all keys into a set */
-    while ((key = lys_getnext(key, snode, NULL, 0)) && (snode->flags & LYS_KEY)) {
-        LY_CHECK_GOTO(ret = ly_set_add(&key_set, (void *)snode, 1, NULL), cleanup);
+    while ((key = lys_getnext(key, snode, NULL, 0)) && (key->flags & LYS_KEY)) {
+        LY_CHECK_GOTO(ret = ly_set_add(&key_set, (void *)key, 1, NULL), cleanup);
     }
 
     LY_LIST_FOR(lyd_child(node), child) {