Do multiple optimizer steps in one training step
What is your question?
Hi!, I'm currently trying to train a GAN with some regularization and I would need to perform multiple optimizer steps with different losses in one single training step.
My intuition tells me to define the optimizer_step function, but I keep finding errors and I don't know what I am doing wrong.

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time

Code
def training_step(self, batch, batch_idx, optimizer_idx):
    real_img = batch
    batch_size = rela_img.shape[0]
    z = torch.randn(batch_size, self.latent_dim).to(self.device)

    # Train Discriminator 
    if optimizer_idx == 0:
        fake_img = self.generator(z)
        fake_pred = self.discriminator(fake_img)
        real_pred = self.discriminator(real_img)
        d_loss = losses.d_loss(real_pred, fake_pred)
        if batch_idx % self.args.d_regularize_every == 0:
            real_img.requires_grad = True
            real_pred = self.discriminator(real_img)
            self.d_reg_loss = losses.d_reg_loss(real_pred, real_img)
        return {'loss': d_loss}

    # Train Generator
    if optimizer_idx == 1:
        fake_img = self.generator(z)
        fake_pred = self.discriminator(fake_img)
        g_loss = losses.g_loss(fake_pred)
        if batch_idx % self.args.g_regularize_every == 0:
             self.g_reg_loss = losses.g_reg_loss(fake_img, z)
        return {'loss':g_loss}

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, *args, **kwargs):
     # Step using d_loss or g_loss
     super().optimizer_step(epoch, batch_idx, optimizer, optimizer_idx, *args, **kwargs)
     if optimizer_idx == 0:
         self.discriminator.zero_grad()
         self.d_reg_loss.backward()
         super().optimizer_step(epoch, batch_idx, optimizer, optimizer_idx, *args, **kwargs)
     if optimizer_idx == 1:
         self.generator.zero_grad()
         self.g_reg_loss.backward()
         super().optimizer_step(epoch, batch_idx, optimizer, optimizer_idx, *args, **kwargs)
What have you tried?
I have tried to directly return the regularization losses instead of the normal losses. But then some at some steps instead of optimizing using the logistic loss I only optimize using the regularization loss. While I would like to always optimize at least one time and then optimize again at some steps.
Is it possible with lightning? Or I am doing something wrong?
What's your environment?

OS: MacOS
Packaging: pip
Version: 0.8.4