diff --git a/neuralmagic/tools/profiler/visualize_trace.py b/neuralmagic/tools/profiler/visualize_trace.py
index b45d0f83fdbf4..8d9ff24916b8a 100644
--- a/neuralmagic/tools/profiler/visualize_trace.py
+++ b/neuralmagic/tools/profiler/visualize_trace.py
@@ -268,7 +268,7 @@ def plot_trace_df(traces_df: pd.DataFrame,
 
     # Setup labels and title
     plt.setp(ax.get_xticklabels(), rotation=90)
-    ax.set_ylabel("time ms")
+    ax.set_ylabel(plot_metric)
     plt.suptitle(plot_title)
 
     plt.savefig(output, bbox_inches='tight')