diff --git a/aidge_core/export_utils/scheduler_export.py b/aidge_core/export_utils/scheduler_export.py index b8a0e1f5dab865b53d390b50d8305d76133e456b..df0b4a385327e4bdccd6fe4de46043d151658dbd 100644 --- a/aidge_core/export_utils/scheduler_export.py +++ b/aidge_core/export_utils/scheduler_export.py @@ -6,7 +6,7 @@ from aidge_core.export_utils import ExportLib, generate_file, copy_file from typing import List, Tuple -def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = None, memory_manager=None, memory_manager_args=None) -> None: +def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = None, memory_manager=None, memory_manager_args=None, labels=False) -> None: graphview = scheduler.graph_view() export_folder = Path().absolute() / export_folder_path @@ -137,11 +137,10 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = if len(outputs_name) != len(outputs_dtype) or len(outputs_name) != len(outputs_size): raise RuntimeError("FATAL: Output args list does not have the same length this is an internal bug.") - if export_lib is not None: - if export_lib.main_jinja_path is not None: - main_jinja_path = export_lib.main_jinja_path - else : - main_jinja_path = str(ROOT / "templates" / "main.jinja") + if export_lib is not None and export_lib.main_jinja_path is not None: + main_jinja_path = export_lib.main_jinja_path + else : + main_jinja_path = str(ROOT / "templates" / "main.jinja") generate_file( str(export_folder / "main.cpp"), @@ -150,7 +149,8 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = inputs_name=inputs_name, outputs_name=outputs_name, outputs_dtype=outputs_dtype, - outputs_size=outputs_size + outputs_size=outputs_size, + labels=labels ) if export_lib is not None: