diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py index b607a19f3fa73f9d99d3509a7918f2a0b7e8b4b1..720208580eda8708f5790c80c995a0fc98565409 100644 --- a/aidge_core/mem_info.py +++ b/aidge_core/mem_info.py @@ -75,8 +75,9 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder mem_planes = mem_manager.get_planes() for node in scheduler.get_static_scheduling(): + node_mem_info = [] if node.type() == "Producer": - mem_info[node] = [] # No meminfo for producer + pass elif node in nodes_at_input: # Input memory management (suppose tensor ends with [:, channel, height, width])) tensor = node.get_operator().get_output(0) @@ -102,19 +103,7 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder "wrap_offset": 0, # No wrapping "wrap_size": 0 # No wrapping }) - mem_info[node] = [{ - "size": plane.size, - "offset": plane.offset, - "stride": plane.stride, - "length": plane.length, - "count": plane.count, - "cont_offset": plane.get_contiguous_offset(), - "cont_size": plane.get_contiguous_size(), - "wrap_offset": plane.get_wrapped_offset(), - "wrap_size": plane.get_wrapped_size() - }] else: - node_mem_info = [] for out_id in range(node.get_nb_outputs()): plane = mem_planes[node][out_id] node_mem_info.append({ @@ -128,7 +117,7 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder "wrap_offset": plane.get_wrapped_offset(), "wrap_size": plane.get_wrapped_size() }) - mem_info[node] = node_mem_info + mem_info[node] = node_mem_info return mem_size, mem_info