diff --git a/_unittests/ut_helpers/test_dot_helper.py b/_unittests/ut_helpers/test_dot_helper.py
index 2e591852..5b250bc5 100644
--- a/_unittests/ut_helpers/test_dot_helper.py
+++ b/_unittests/ut_helpers/test_dot_helper.py
@@ -54,7 +54,7 @@ def test_custom_doc_kernels_layer_normalization(self):
               I_2 -> LayerNormalization_3 [label="FLOAT16(d)"];
               LayerNormalization_3 -> Add_4 [label="FLOAT16(b,c,d)"];
               I_1 -> Add_4 [label="FLOAT16(d)"];
-              O_5 [label="Z\\nFLOAT16(d)", fillcolor="#aaaaee"];
+              O_5 [label="Z\\nFLOAT16(b,c,d)", fillcolor="#aaaaee"];
               Add_4 -> O_5;
             }
             """
diff --git a/onnx_diagnostic/helpers/dot_helper.py b/onnx_diagnostic/helpers/dot_helper.py
index 328446cd..a1d6d1a9 100644
--- a/onnx_diagnostic/helpers/dot_helper.py
+++ b/onnx_diagnostic/helpers/dot_helper.py
@@ -201,7 +201,7 @@ def _mkn(obj: object) -> int:
     for out in outputs:
         if not out.name:
             continue
-        lab = _make_edge_label(inp)
+        lab = _make_edge_label(out)
         rows.append(f'  O_{_mkn(out)} [label="{out.name}\\n{lab}", fillcolor="#aaaaee"];')
         edge = name_to_ids[out.name], f"O_{_mkn(out)}"
         rows.append(f"  {edge[0]} -> {edge[1]};")