diff --git a/aidge_core/export_utils/scheduler_export.py b/aidge_core/export_utils/scheduler_export.py
index 310e7e6508834b127a8af0f7dbf5a010b546141b..98c78b087eb539cc6b229948bc81180e6a152fc8 100644
--- a/aidge_core/export_utils/scheduler_export.py
+++ b/aidge_core/export_utils/scheduler_export.py
@@ -172,7 +172,8 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib =
         inputs_name=inputs_name,
         inputs_dtype=inputs_dtype,
         outputs_name=outputs_name,
-        outputs_dtype=outputs_dtype
+        outputs_dtype=outputs_dtype,
+        test_mode=test_mode
     )
 
     forward_header_template = str(ROOT / "templates" / "forward_header.jinja")
@@ -188,7 +189,8 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib =
         inputs_name=inputs_name,
         inputs_dtype=inputs_dtype,
         outputs_name=outputs_name,
-        outputs_dtype=outputs_dtype
+        outputs_dtype=outputs_dtype,
+        test_mode=test_mode
     )
 
     if len(outputs_name) != len(outputs_dtype) or len(outputs_name) != len(outputs_size):