Repeated calls to generator and discriminator's forward in GAN tutorial

[jhauret]

📚 Documentation

In the training_step of the GAN(L.LightningModule) the generator and discriminator forward are called several times on the same input. Obviously this slows down the training because more computation is required. I wonder if we could just reuse the results of the first call. After all, the toggle_optimiser/untoggle_optimiser functions should make it safe, right?

For the generator:

• First call: to log images self.generated_imgs = self(z)
• Second call: Inside the generator optimization self.discriminator(self(z))
• Third call: Inside the discriminator optimization self.discriminator(self(z).detach())

For the discriminator:

• First call: Inside the generator optimization self.discriminator(self(z))
• Second call: Inside the discriminator optimization self.discriminator(self(z).detach())

cc @borda