How to get peak memory of solution

[haoming-codes]

from checkmate.core.solvers.gurobi_solver import solve_ilp_gurobi as solver
from checkmate.core.graph_builder import gen_linear_graph
g = gen_linear_graph(5)
solver(g, 3, approx=False)
print(sched_result.solver_budget)
print(sched_result.schedule_aux_data.peak_ram)

Here, the peak ram == sum of all tensor size, even though the schedule sched_result.schedule clearly is the peak memory = 3 solution that rematerializes every forward node during the backward pass. May I know how to get the correct peak memory from a ScheduledResult? Thank you!