Stability-AI/generative-models · error · NotImplementedError

Unknown optimizer {optimizer_idx}

Error message

Unknown optimizer {optimizer_idx}

What it means

AutoencodingEngine's inner_training_step handles two optimizer indices: 0 for the generator/autoencoder and 1 for the discriminator. Any other optimizer_idx (Lightning's optimize loop index) reaches the else branch and raises this NotImplementedError. This normally indicates the Lightning configure_optimizers returned more optimizers than the training step supports.

Source

Thrown at sgm/models/autoencoder.py:279

            self.log(
                "loss",
                aeloss.mean().detach(),
                prog_bar=True,
                logger=False,
                on_epoch=False,
                on_step=True,
            )
            return aeloss
        elif optimizer_idx == 1:
            # discriminator
            discloss, log_dict_disc = self.loss(x, xrec, **extra_info)
            # -> discriminator always needs to return a tuple
            self.log_dict(
                log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True
            )
            return discloss
        else:
            raise NotImplementedError(f"Unknown optimizer {optimizer_idx}")

    def training_step(self, batch: dict, batch_idx: int):
        opts = self.optimizers()
        if not isinstance(opts, list):
            # Non-adversarial case
            opts = [opts]
        optimizer_idx = batch_idx % len(opts)
        if self.global_step < self.disc_start_iter:
            optimizer_idx = 0
        opt = opts[optimizer_idx]
        opt.zero_grad()
        with opt.toggle_model():
            loss = self.inner_training_step(
                batch, batch_idx, optimizer_idx=optimizer_idx
            )
            self.manual_backward(loss)
        opt.step()

View on GitHub (pinned to e8cd657656)

Solutions

1. Ensure configure_optimizers returns exactly two optimizers (autoencoder and discriminator) when using adversarial training.
2. Pin/check the PyTorch Lightning version for changes in multi-optimizer training_step invocation semantics.
3. If adding a third training objective, extend inner_training_step to handle the new optimizer_idx instead of only 0/1.

opts = model.configure_optimizers()
if isinstance(opts, tuple):
    opts = opts[0]
assert len(opts) == 2, "AutoencodingEngine supports exactly 2 optimizers (gen + disc)"

try:
    trainer.fit(model, datamodule=dm)
except NotImplementedError as e:
    if "Unknown optimizer" in str(e):
        raise RuntimeError("configure_optimizers must return exactly 2 optimizers") from e

When it happens

Trigger: Training the autoencoder with a Lightning version or subclass whose training loop calls inner_training_step with an optimizer_idx other than 0/1, or overriding configure_optimizers to return 3+ optimizers.

Common situations: Migrating across PyTorch Lightning versions where automatic optimization invokes training_step per optimizer; custom autoencoder subclasses returning extra optimizers; mis-wired multiple-optimizer configs.

Related errors

• Unknown optimizer_idx {optimizer_idx}
• Checkpoint path is deprecated, use `checkpoint_egnine` inste
• Did not find parameters for pattern {pattern_}

AI-assisted analysis of Stability-AI/generative-models@e8cd657656 (2026-08-29). Data as JSON: /api/errors/83288679c2ecfaae. Report an issue: GitHub.