Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -252,11 +252,20 @@ move_label_to_readout = tfgnn.keras.layers.StructuredReadoutIntoFeature(
 
 feature_processors = [process_features, add_readout, move_label_to_readout]
 
+# Function to create the full model with feature processing
+def create_full_model(graph_tensor_spec: tfgnn.GraphTensorSpec):
+    input_graph = tf.keras.layers.Input(type_spec=graph_tensor_spec)
+    graph = input_graph
+    for processor in feature_processors:
+        graph = processor(graph)
+    output_graph = model_fn(graph.spec)(graph)
+    return tf.keras.Model(input_graph, output_graph)
+
 # Run training
 st.write("Training the model...")
 runner.run(
     task=task,
-    model_fn=model_fn,
+    model_fn=create_full_model,
     trainer=trainer,
     optimizer_fn=optimizer_fn,
     epochs=epochs,