Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -171,8 +171,6 @@ def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
             normalization_type="layer",
             next_state_type="residual",
         )(graph)
-    graph = tfgnn.keras.layers.Readout(
-        node_set_name="paper", feature_name="feat")(graph)
     return tf.keras.Model(inputs, graph)
 
 # Check for TPU/ GPU and set strategy
@@ -260,6 +258,9 @@ def create_full_model(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = input_graph
     for processor in feature_processors:
         graph = processor(graph)
+    # Ensure the Readout layer is correctly added after feature processing
+    readout = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="feat")
+    graph = readout(graph)
     output_graph = model_fn(graph.spec)(graph)
     return tf.keras.Model(input_graph, output_graph)