diff --git a/aidge_core/export_utils/node_export.py b/aidge_core/export_utils/node_export.py index 5eba1647cd0c45cb0eba48ea83bb36534870e3d8..581e9070b409621c845e69c5316c3f6597237434 100644 --- a/aidge_core/export_utils/node_export.py +++ b/aidge_core/export_utils/node_export.py @@ -113,7 +113,7 @@ class ExportNode(ABC): """ @abstractmethod - def __init__(self, aidge_node: aidge_core.Node, mem_info: List[dict], is_input: bool, is_output: bool) -> None: + def __init__(self, aidge_node: aidge_core.Node, mem_info: List[dict]=None, is_input: bool=False, is_output: bool=False) -> None: """Create ExportNode and retieve attriubtes from ``aidge_node``: """ @@ -159,7 +159,7 @@ class ExportNode(ABC): # Producer don't have meminfo # TODO: document this attribute # true if node have meminfo else false - self.attributes["meminfo"] = self.node.type() != "Producer" + self.attributes["meminfo"] = mem_info or self.node.type() != "Producer" if self.attributes["meminfo"]: self.attributes["mem_info_size"] = [None] * self.attributes["nb_out"] self.attributes["mem_info_offset"] = [None] * self.attributes["nb_out"]