diff --git a/src/tree_schema_compile.c b/src/tree_schema_compile.c
index cd445b2..a81609f 100644
--- a/src/tree_schema_compile.c
+++ b/src/tree_schema_compile.c
@@ -576,7 +576,7 @@
 
     ext->insubstmt = ext_p->insubstmt;
     ext->insubstmt_index = ext_p->insubstmt_index;
-    ext->module = ctx->mod_def;
+    ext->module = ctx->mod;
     ext->parent = parent;
     ext->parent_type = parent_type;