Update pages/15_Graphs.py

pages/15_Graphs.py

@@ -13,9 +13,9 @@ os.environ['TF_USE_LEGACY_KERAS'] = '1'
 def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec)
 
-    # Encode input features
+    # Encode input features to match the required output shape of 128
     graph = tfgnn.keras.layers.MapFeatures(
-        node_sets_fn=lambda node_set, node_set_name: tf.keras.layers.Dense(64)(node_set['features'])
+        node_sets_fn=lambda node_set, node_set_name: tf.keras.layers.Dense(128)(node_set['features'])
     )(graph)
 
     # For each round of message passing...
@@ -39,7 +39,7 @@ def create_sample_graph():
     graph = nx.gnm_random_graph(num_nodes, num_edges, directed=True)
 
     # Create a GraphTensor
-    node_features = tf.random.normal((num_nodes, 64))
+    node_features = tf.random.normal((num_nodes, 128))  # Match the dense layer output
     edge_features = tf.random.normal((num_edges, 32))
 
     graph_tensor = tfgnn.GraphTensor.from_pieces(