From 708b40967043008d13a793496f27b6d6105b78af Mon Sep 17 00:00:00 2001
From: Dobson <bdobson@ic.ac.uk>
Date: Fri, 15 Mar 2024 16:41:08 +0000
Subject: [PATCH] Update graph_utilities.py

---
 swmmanywhere/graph_utilities.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/swmmanywhere/graph_utilities.py b/swmmanywhere/graph_utilities.py
index ccad72a4..3c1cef45 100644
--- a/swmmanywhere/graph_utilities.py
+++ b/swmmanywhere/graph_utilities.py
@@ -1043,7 +1043,7 @@ def __call__(self,
             G (nx.Graph): A graph
         """
         G = G.copy()
-        surface_elevations = {n : d['surface_elevation'] for n, d in G.nodes(data=True)}
+        surface_elevations = nx.get_node_attributes(G, 'surface_elevation')
         topological_order = list(nx.topological_sort(G))
         chamber_floor = {}
         edge_diams: dict[tuple[Hashable,Hashable,int],float] = {}