How to get the output of train_on_batch in Keras?

I am trying to train a GAN. I would like to be able to capture the predicted output of my generator when I call train_on_batch. This way I could use that output to train the discriminator. Right now I have to call predict on my generator to generate new training samples. This is less efficient. I saw a way to do it with an older version by adding a metric but it appears they removed it from the latest version. I am using Keras included with Tensorflow 2.4.1.