Update heterograph.py

Protect temporary changes from persisting if an error occurs during the yield block.
dmlc · Oct 27, 2023 · 26564f4 · 26564f4
1 parent a93b6fe
commit 26564f4
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
@@ -6061,9 +6061,11 @@ def local_scope(self):
         old_eframes = self._edge_frames
         self._node_frames = [fr.clone() for fr in self._node_frames]
         self._edge_frames = [fr.clone() for fr in self._edge_frames]
-        yield
-        self._node_frames = old_nframes
-        self._edge_frames = old_eframes
+        try:
+            yield
+        finally:
+            self._node_frames = old_nframes
+            self._edge_frames = old_eframes
 
     def formats(self, formats=None):
         r"""Get a cloned graph with the specified allowed sparse format(s) or