fix hamilton graph network example

google-deepmind · Jan 7, 2023 · e5b1659 · e5b1659
1 parent 51f5990
commit e5b1659
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jraph/examples/hamiltonian_graph_network.py b/jraph/examples/hamiltonian_graph_network.py
@@ -199,7 +199,8 @@ def set_system_state(
     position: np.ndarray,
     momentum: np.ndarray) -> jraph.GraphsTuple:
   """Sets the non-static parameters of the graph (momentum, position)."""
-  nodes = static_graph.nodes.copy(position=position, momentum=momentum)
+  nodes = static_graph.nodes.set("position", position)
+  nodes = nodes.set("momentum", momentum)
   return static_graph._replace(nodes=nodes)