diff --git a/src/itwinai/cli.py b/src/itwinai/cli.py
index afd6269a..286a99ad 100644
--- a/src/itwinai/cli.py
+++ b/src/itwinai/cli.py
@@ -157,7 +157,7 @@ def generate_communication_plot(
 
     from itwinai.scalability import convert_matching_files_to_dataframe
     from itwinai.torch.profiling.communication_plot import (
-        create_stacked_plot,
+        communication_overhead_stacked_bar_plot,
         get_comp_fraction_full_array,
     )
 
@@ -185,7 +185,7 @@ def generate_communication_plot(
     strategies = sorted(communication_df["strategy"].unique())
     gpu_numbers = sorted(communication_df["num_gpus"].unique(), key=lambda x: int(x))
 
-    fig, _ = create_stacked_plot(values, strategies, gpu_numbers)
+    fig, _ = communication_overhead_stacked_bar_plot(values, strategies, gpu_numbers)
 
     # TODO: set these dynamically?
     fig.set_figwidth(8)
diff --git a/src/itwinai/torch/profiling/communication_plot.py b/src/itwinai/torch/profiling/communication_plot.py
index 0797119c..3d9a704e 100644
--- a/src/itwinai/torch/profiling/communication_plot.py
+++ b/src/itwinai/torch/profiling/communication_plot.py
@@ -64,7 +64,7 @@ def calculate_comp_and_comm_time(df: pd.DataFrame) -> Tuple[float, float]:
     return comp_time, comm_time
 
 
-def create_stacked_plot(
+def communication_overhead_stacked_bar_plot(
     values: np.ndarray, strategy_labels: List, gpu_numbers: List
 ) -> Tuple[Any, Any]:
     """Creates a stacked plot showing values from 0 to 1, where the given value
@@ -145,10 +145,9 @@ def create_stacked_plot(
 def get_comp_fraction_full_array(
     df: pd.DataFrame, print_table: bool = False
 ) -> np.ndarray:
-    """Creates a MxN NumPy array where M is the number of strategies
-    and N is the number of GPU configurations. The strategies are sorted
-    alphabetically and the GPU configurations are sorted in ascending number
-    of GPUs.
+    """Creates a MxN NumPy array where M is the number of strategies and N is the
+    number of GPU configurations. The strategies are sorted alphabetically and the GPU
+    configurations are sorted in ascending number of GPUs.
     """
     unique_num_gpus = sorted(df["num_gpus"].unique(), key=lambda x: int(x))
     unique_strategies = sorted(df["strategy"].unique())
@@ -165,8 +164,7 @@ def get_comp_fraction_full_array(
 
             row_string = f"{strategy:>12} | {num_gpus:>10}"
 
-            # Allows asymmetric testing, i.e. not testing all num gpus and all
-            # strategies together
+            # Allows some strategies or num GPUs to not be included
             if len(filtered_df) == 0:
                 comp_time, comm_time = np.NaN, np.NaN
                 strategy_values.append(np.NaN)