diff --git a/src/query/procedure/module.cpp b/src/query/procedure/module.cpp
index 69582a22e..dc28da21c 100644
--- a/src/query/procedure/module.cpp
+++ b/src/query/procedure/module.cpp
@@ -22,8 +22,8 @@ std::optional<Module> LoadModuleFromSharedLibrary(std::filesystem::path path) {
     return std::nullopt;
   }
   // Get required mgp_main
-  module.main_fn = reinterpret_cast<void (*)(const mgp_list *,
-                                             const mgp_graph *, mgp_result *)>(
+  module.main_fn = reinterpret_cast<void (*)(
+      const mgp_list *, const mgp_graph *, mgp_result *, mgp_memory *)>(
       dlsym(module.handle, "mgp_main"));
   const char *error = dlerror();
   if (!module.main_fn || error) {
diff --git a/src/query/procedure/module.hpp b/src/query/procedure/module.hpp
index 82cd51da5..9518ed287 100644
--- a/src/query/procedure/module.hpp
+++ b/src/query/procedure/module.hpp
@@ -13,6 +13,7 @@
 
 struct mgp_graph;
 struct mgp_list;
+struct mgp_memory;
 struct mgp_result;
 
 namespace query::procedure {
@@ -23,7 +24,8 @@ struct Module final {
   /// System handle to shared library.
   void *handle;
   /// Entry-point for module's custom procedure.
-  std::function<void(const mgp_list *, const mgp_graph *, mgp_result *)>
+  std::function<void(const mgp_list *, const mgp_graph *, mgp_result *,
+                     mgp_memory *)>
       main_fn;
   /// Optional initialization function called on module load.
   std::function<int()> init_fn;