diff --git a/aidge_core/export_utils/scheduler_export.py b/aidge_core/export_utils/scheduler_export.py
index e7b426b7cc17274b7bd8ed998bd4a4549ef03490..b39512b19900b99743ee4aea236d736eb37f2f38 100644
--- a/aidge_core/export_utils/scheduler_export.py
+++ b/aidge_core/export_utils/scheduler_export.py
@@ -157,7 +157,11 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib =
         func_name=func_name,
         headers=set(list_configs),
         actions=list_actions,
-        mem_ctype=inputs_dtype[0],  # Legacy behavior ...
+        # Note: Graph may not have inputs, so we need to check with output
+        # In the future, we should remove this has it is not compatible
+        # with a mix precision approach.
+        mem_ctype=outputs_dtype[0],  # Legacy behavior ...
+        mem_section=export_lib.mem_section,
         peak_mem=peak_mem,
         inputs_name=inputs_name,
         inputs_dtype=inputs_dtype,
diff --git a/aidge_core/export_utils/templates/main.jinja b/aidge_core/export_utils/templates/main.jinja
index d0c22719a5d9b9eaa15d3ef9ef86307a060b54be..697a97b53ef05eedddf5ac66f262581475e655c6 100644
--- a/aidge_core/export_utils/templates/main.jinja
+++ b/aidge_core/export_utils/templates/main.jinja
@@ -26,7 +26,7 @@ int main()
     {% endfor %}
 
     // Call the forward function
-    {{ func_name }}({{ inputs_name|join(", ") }}, &{{ outputs_name|join(", &") }});
+    {{ func_name }}({{ inputs_name|join(", ") }}{% if inputs_name %}, {% endif %}&{{ outputs_name|join(", &") }});
 
     // Print the results of each output
     {%- for o in range(outputs_name | length) %}