code: def train():
      with tf.GradientTape() as tape:
          predictions = model([x, a], training=True)
          loss = loss_fn(y[mask_tr], predictions[mask_tr])
          loss += sum(model.losses)
      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(
          zip(gradients, model.trainable_variables)
      )
      return loss
  
file_path: spektral/examples/node_prediction/citation_gat_custom.py
imports: import tensorflow
  from tensorflow.keras.layers import Dropout, Input
  from tensorflow.keras.losses import CategoricalCrossentropy
  from tensorflow.keras.metrics import categorical_accuracy
  from tensorflow.keras.models import Model
  from tensorflow.keras.optimizers import Adam
  from tensorflow.keras.regularizers import l2
  from spektral.datasets.citation import Cora
  from spektral.layers import GATConv
  from spektral.transforms import AdjToSpTensor, LayerPreprocess
  from spektral.utils import tic, toc
label_desc: Use image node_id 1 for calling a global function with example usage: train() and returns: loss