code: def train_step(inputs, target):
      with tf.GradientTape() as tape:
          predictions = model(inputs, training=True)
          loss = loss_fn(target, predictions) + sum(model.losses)
      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(
          zip(gradients, model.trainable_variables)
      )
      acc = tf.reduce_mean(categorical_accuracy(target, predictions))
      return loss, acc
  
file_path: spektral/examples/graph_prediction/general_gnn.py
imports: import numpy
  import tensorflow
  from tensorflow.keras.losses import CategoricalCrossentropy
  from tensorflow.keras.metrics import categorical_accuracy
  from tensorflow.keras.optimizers import Adam
  from spektral.data import DisjointLoader
  from spektral.datasets import TUDataset
  from spektral.models import GeneralGNN
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(inputs, target) and returns: loss,  acc