code: def train_model(model, train_x, train_y, save_weights=False):
      """Train the model on given data."""
      optimizer = dp_optimizer_vectorized.VectorizedDPSGD(
          l2_norm_clip=FLAGS.l2_norm_clip,
          noise_multiplier=FLAGS.noise_multiplier,
          num_microbatches=FLAGS.microbatches,
          learning_rate=FLAGS.learning_rate,
      )
  
      loss = tf.keras.losses.CategoricalCrossentropy(
          from_logits=True, reduction=tf.losses.Reduction.NONE
      )
  
      # Compile model with Keras
      model.compile(
          optimizer=optimizer, loss=loss, metrics=["accuracy"]
      )
  
      if save_weights:
          wts = model.get_weights()
          np.save("save_model", wts)
          model.set_weights(wts)
          return model
  
      if FLAGS.load_weights:  # load preset weights
          wts = np.load("save_model.npy", allow_pickle=True).tolist()
          model.set_weights(wts)
  
      # Train model with Keras
      model.fit(
          train_x,
          train_y,
          epochs=FLAGS.epochs,
          validation_data=(train_x, train_y),
          batch_size=FLAGS.batch_size,
          verbose=0,
      )
      return model
  
file_path: privacy/research/audit_2020/fmnist_audit.py
imports: import numpy
  import tensorflow.compat.v1
  from tensorflow_privacy.privacy.analysis.rdp_accountant import compute_rdp
  from tensorflow_privacy.privacy.analysis.rdp_accountant import get_privacy_spent
  from tensorflow_privacy.privacy.optimizers import dp_optimizer_vectorized
  from absl import app
  from absl import flags
  import audit
label_desc: Use image node_id 3 for calling a global function with example usage: train_model(model, train_x, train_y, save_weights) and returns: model, model