Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -150,7 +150,7 @@ def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
     if node_set_name == "institution":
         return tf.keras.layers.Embedding(6_500, 16)(node_set["hashed_id"])
     if node_set_name == "paper":
-        return tf.keras.layers.Dense(node_state_dim, "relu")(node_set["feat"])
+        return tf.keras.layers.Dense(node_state_dim, activation="relu")(node_set["feat"])
     if node_set_name == "author":
         return node_set["empty_state"]
     raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
@@ -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)
     return tf.keras.Model(inputs, graph)
 
 # Check for TPU/ GPU and set strategy
@@ -191,7 +193,7 @@ else:
     train_padding = None
     valid_padding = None
 
-st.write(f"Found {strategy.num_replicas_in_sync} replicas in sync")
+st.write(f"Found {strategy.num_replicas in_sync} replicas in sync")
 
 # Define task
 st.write("Defining the task...")