diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py
index 53ff73cef2d8d2f17428b5fc3f9917e94c7dce8a..dc96b49cd12a1c10bc599057e2c5feef821b5cff 100644
--- a/aidge_core/mem_info.py
+++ b/aidge_core/mem_info.py
@@ -45,17 +45,17 @@ def compute_default_mem_info(scheduler: aidge_core.Scheduler) -> Tuple[int, List
     return mem_size, mem_info
 
 # TODO remove
-# def _gnuplot_installed():
-#     try:
-#         # Run gnuplot with the --version flag and capture the output
-#         subprocess.run(["gnuplot", "--version"])
-#         return True
-#     except FileNotFoundError:
-#         aidge_core.Log.warn("Gnuplot is not installed.")
-#         return False
-#     except subprocess.CalledProcessError:
-#         aidge_core.Log.warn("Gnuplot command found but failed to run.")
-#         return False
+def _gnuplot_installed():
+    try:
+        # Run gnuplot with the --version flag and capture the output
+        subprocess.run(["gnuplot", "--version"])
+        return True
+    except FileNotFoundError:
+        aidge_core.Log.warn("Gnuplot is not installed.")
+        return False
+    except subprocess.CalledProcessError:
+        aidge_core.Log.warn("Gnuplot command found but failed to run.")
+        return False
 
 def log_meminfo(mem_manager:aidge_core.MemoryManager, path: Path, diplay_names:bool):
     """Generate a graph representing the memory allocation of each ouputs.
@@ -104,12 +104,17 @@ def log_meminfo(mem_manager:aidge_core.MemoryManager, path: Path, diplay_names:b
                 plt.text(x_end,y_end, node.name(), rotation=45)
         color_id += 1
 
+    plt.xlim(0, max_lifetime + 1)
+    plt.ylim(0, peak_usage)
     plt.axhline(y=peak_usage, color='red', linestyle='--')
     plt.text(0, peak_usage, f'Peak usage = {peak_usage} KWords', color='red')
     plt.xlabel("Time")
     plt.ylabel("Memory usage (KWords)")
     plt.title("Memory Usage Over Time")
     plt.grid(True)
+    ax = plt.gca()
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
     plt.savefig(path)
     plt.close()
     aidge_core.Log.notice(f"Generated memory management info at: {path}")
@@ -153,18 +158,18 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder
     if stats_folder is not None:
         log_meminfo(mem_manager, Path(stats_folder) / "memory_info.png", display_names)
         # TODO remove
-        # if _gnuplot_installed():
-        #     # Use gnuplot to generate the log
-        #     os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True)
-        #     mem_manager.log("memory_info")
-        #     os.chmod("memory_info_plot.gnu", 0o777)
-        #     os.system("./memory_info_plot.gnu")
-        #     shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info"))
-        #     shutil.move("memory_info_plot.png", str(
-        #         Path(stats_folder) / "graph" / "memory_info_plot.png"))
-        #     os.remove("memory_info_plot.gnu")
-        # else:
-        #     aidge_core.Log.warn("Warning: gnuplot is not installed, could not generate stat folder.")
+        if _gnuplot_installed():
+            # Use gnuplot to generate the log
+            os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True)
+            mem_manager.log("memory_info")
+            os.chmod("memory_info_plot.gnu", 0o777)
+            os.system("./memory_info_plot.gnu")
+            shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info"))
+            shutil.move("memory_info_plot.png", str(
+                Path(stats_folder) / "graph" / "memory_info_plot.png"))
+            os.remove("memory_info_plot.gnu")
+        else:
+            aidge_core.Log.warn("Warning: gnuplot is not installed, could not generate stat folder.")
 
     # In the export, we currently use an unified memory buffer whose size
     # is determined by the memory peak usage