Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -171,6 +171,8 @@ 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)
+    graph = tf.keras.layers.Dense(349, activation="softmax")(graph)
     return tf.keras.Model(inputs, graph)
 
 # Check for TPU/ GPU and set strategy
@@ -258,10 +260,8 @@ 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)
+    graph = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="feat")(graph)
+    output_graph = model_fn(graph_tensor_spec)(graph)
     return tf.keras.Model(input_graph, output_graph)
 
 # Run training