.

utils.py

@@ -146,6 +146,7 @@ def plot_memory_breakdown(
         seq_len, mbs, batch_accum, tp, pp, dp, zero_stage,
         tie_word_embeddings
     )
+    memory_usage_peak_tbi = results["Aggregates"]["Peak Memory (TBI)"]
     
     # Create figure for components plot
     plt.close('all')
@@ -261,4 +262,4 @@ def plot_memory_breakdown(
     # Add legend below the plot
     # plt.subplots_adjust(bottom=0.8)
     ax2.legend(loc='lower center', bbox_to_anchor=(0.5, -1.5), ncol=3)
-    return fig1, fig2
+    return fig1, fig2, memory_usage_peak_tbi