accum-batches can't be used with autoencoder

[Jourdelune]

When we set "model_type": "autoencoder" in the configuration of the model, the argument accum-batches can't be used because automatic_optimization is disabled.

MisconfigurationException: Automatic gradient accumulation is not supported for manual optimization. Remove `Trainer(accumulate_grad_batches=32)` or switch to automatic optimization.

To fix the issue, I have written the code to accumulate the batch in the AutoencoderTrainingWrapper class.

if use_disc:
            loss, losses = self.losses_disc(loss_info)

            log_dict["train/disc_lr"] = opt_disc.param_groups[0]["lr"]

            # Gradient accumulation
            if self.gradient_accumulation is not None:
                loss = loss / self.gradient_accumulation

            self.manual_backward(loss)
            if self.clip_grad_norm > 0.0:
                torch.nn.utils.clip_grad_norm_(
                    self.discriminator.parameters(), self.clip_grad_norm
                )

            # accumulate gradients of N batches
            if (batch_idx + 1) % self.gradient_accumulation == 0:
                opt_disc.step()

                if sched_disc is not None:
                    # sched step every step
                    sched_disc.step()

                opt_disc.zero_grad()

        # Train the generator
        else:
            loss, losses = self.losses_gen(loss_info)

            if self.use_ema:
                self.autoencoder_ema.update()

            log_dict["train/loss"] = loss.detach().item()
            log_dict["train/latent_std"] = latents.std().detach().item()
            log_dict["train/data_std"] = data_std.detach().item()
            log_dict["train/gen_lr"] = opt_gen.param_groups[0]["lr"]

            # Gradient accumulation
            if self.gradient_accumulation is not None:
                loss = loss / self.gradient_accumulation

            self.manual_backward(loss)
            if self.clip_grad_norm > 0.0:
                torch.nn.utils.clip_grad_norm_(
                    self.autoencoder.parameters(), self.clip_grad_norm
                )

            # accumulate gradients of N batches
            if (batch_idx + 1) % self.gradient_accumulation == 0:
                opt_gen.step()

                if sched_gen is not None:
                    # scheduler step every step
                    sched_gen.step()

                opt_gen.zero_grad()