Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -155,10 +155,9 @@ def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
         return node_set["empty_state"]
     raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
 
-# Define the GNN model function
 def gnn_model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
-    graph = inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
-    graph = tfgnn.keras.layers.MapFeatures(node_sets_fn=set_initial_node_states)(graph)
+    inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
+    graph = tfgnn.keras.layers.MapFeatures(node_sets_fn=set_initial_node_states)(inputs)
     for i in range(num_graph_updates):
         graph = mt_albis.MtAlbisGraphUpdate(
             units=node_state_dim,
@@ -173,7 +172,6 @@ def gnn_model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
         )(graph)
     return tf.keras.Model(inputs, graph)
 
-# Define the complete model function
 def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     gnn_model = gnn_model_fn(graph_tensor_spec)
     inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)