diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py index 8946c4dbdbab23b29f6c687db2d12026b15343c5..5e36eb3ee56165b9ddd2422a8e0607f79871afe8 100644 --- a/aidge_core/mem_info.py +++ b/aidge_core/mem_info.py @@ -144,19 +144,18 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder }) else: for out_id in range(node.get_nb_outputs()): - if node in mem_planes: - plane = mem_planes[node][out_id] - node_mem_info.append({ - "size": plane.size, - "offset": plane.get_contiguous_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() - }) + plane = mem_planes[node][out_id] + node_mem_info.append({ + "size": plane.size, + "offset": plane.get_contiguous_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() + }) mem_info[node] = node_mem_info return mem_size, mem_info