code: def train(model, adj, features, labels, idx_train, idx_val, idx_test):
      """Train gnn model."""
      loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
          from_logits=True
      )
      best_val_acc = 0.0
  
      if FLAGS.optimizer == "adam":
          optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.lr)
      elif FLAGS.optimizer == "sgd":
          optimizer = tf.keras.optimizers.SGD(learning_rate=FLAGS.lr)
  
      inputs = (features, adj)
      for epoch in range(FLAGS.epochs):
          epoch_start_time = time.time()
  
          with tf.GradientTape() as tape:
              output = model(inputs, training=True)
              train_loss = loss_fn(labels[idx_train], output[idx_train])
              # L2 regularization
              for weight in model.trainable_weights:
                  train_loss += FLAGS.weight_decay * tf.nn.l2_loss(
                      weight
                  )
  
          gradients = tape.gradient(
              train_loss, model.trainable_variables
          )
          optimizer.apply_gradients(
              zip(gradients, model.trainable_variables)
          )
  
          train_acc = cal_acc(labels[idx_train], output[idx_train])
  
          # Evaluate
          output = model(inputs, training=False)
          val_loss = loss_fn(labels[idx_val], output[idx_val])
          val_acc = cal_acc(labels[idx_val], output[idx_val])
  
          if FLAGS.save_best_val:
              if val_acc >= best_val_acc:
                  best_val_acc = val_acc
                  model.save(FLAGS.save_dir)
  
          print(
              "[%03d/%03d] %.2f sec(s) Train Acc: %.3f Loss: %.6f | Val Acc: %.3f loss: %.6f"
              % (
                  epoch + 1,
                  FLAGS.epochs,
                  time.time() - epoch_start_time,
                  train_acc,
                  train_loss,
                  val_acc,
                  val_loss,
              )
          )
  
      if not FLAGS.save_best_val:
          model.save(FLAGS.save_dir)
      print("Start Predicting...")
      model = tf.keras.models.load_model(FLAGS.save_dir)
      output = model(inputs, training=False)
      test_acc = cal_acc(labels[idx_test], output[idx_test])
      print("***Test Accuracy: %.3f***" % test_acc)
  
imports: import time
  from absl import app
  from absl import flags
  import tensorflow
  from utils import load_dataset, build_model, cal_acc
label_desc: Use image node_id 1 for calling a global function with example usage: train(model, adj, features, labels, idx_train, idx_val, idx_test) without return types