diff --git a/src/weather_model_graphs/visualise/plot_2d.py b/src/weather_model_graphs/visualise/plot_2d.py
index 9cf28ed..e4f3ec4 100644
--- a/src/weather_model_graphs/visualise/plot_2d.py
+++ b/src/weather_model_graphs/visualise/plot_2d.py
@@ -146,7 +146,7 @@ def nx_draw_with_pos_and_attr(
             graph, edge_color_attr, component="edges"
         )
 
-        if "cmap" not in kwargs:
+        if "edge_cmap" not in kwargs:
             if "discrete_labels" in edge_attr_vals:
                 kwargs["edge_cmap"] = plt.get_cmap("tab20")
             else: