Как вычислить градиенты с весами в tf.keras и зарегистрировать то же самое для мониторинга градиентов на тензорной доске?

На самом деле я работаю над базовыми GAN, используя TF keras, здесь я использую метод train_on_batch для обучения генератора и дискриминатора, в качестве альтернативы, который не имеет параметра обратных вызовов для записи журналов тензорной доски, например >fit метод модели keras. Теперь я хочу писать журналы модели во время обучения для отслеживания весов и градиентов на тензорной доске.

Часть обучающего кода выглядит следующим образом:

def train(g_model, d_model, gan_model, dataset, latent_dim, seed, n_epochs=100, n_batch=128):
  bat_per_epo = int(dataset.shape[0] / n_batch)
  half_batch = int(n_batch / 2)

  for i in range(n_epochs):
    for j in range(bat_per_epo):
      # Training discriminator with real images
      X_real, y_real = generate_real_samples(dataset, half_batch)
      d_loss1, _ = d_model.train_on_batch(X_real, y_real * .9) # Label Smoothing

      # Training discriminator with fake images
      X_fake, y_fake = generate_fake_samples(g_model, latent_dim, half_batch)
      d_loss2, _ = d_model.train_on_batch(X_fake, y_fake + .1) # Label Smoothing

      # Training generator with latent points
      X_gan = generate_latent_points(latent_dim, n_batch)
      y_gan = ones((n_batch, 1))

      g_loss = gan_model.train_on_batch(X_gan, y_gan)

      if not j%10:
        print('>%d, %d/%d, d1=%.3f, d2=%.3f g=%.3f' %(i+1, j+1, bat_per_epo, d_loss1, d_loss2, g_loss))

    display.clear_output(True)
    print('>%d, %d/%d, d1=%.3f, d2=%.3f g=%.3f' %(i+1, j+1, bat_per_epo, d_loss1, d_loss2, g_loss))
    summarize_performance(i, g_model, d_model, dataset, latent_dim, seed)

  display.clear_output(True)
  print('>%d, %d/%d, d1=%.3f, d2=%.3f g=%.3f' %(i+1, j+1, bat_per_epo, d_loss1, d_loss2, g_loss))
  summarize_performance(i, g_model, d_model, dataset, latent_dim, seed)

Некоторые из способов вычисления градиентов, которые я нашел здесь,

1. Как получить градиенты в keras?
2. Получение градиента вывода модели с весами с использованием Keras

Но я смущен этим и тем, как регистрировать градиенты без опции обратных вызовов. Может кто-нибудь, пожалуйста, помогите мне в этом?