Lightning-AI/pytorch-lightning · error · MisconfigurationException

In automatic_optimization, when `training_step` returns a di

Error message

In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present

What it means

Raised by ClosureResult.from_training_step_output in automatic optimization when training_step returns a dict/Mapping that either lacks the 'loss' key or maps it to None. In automatic optimization Lightning needs a loss tensor to call backward on, so a dict without 'loss' cannot be processed.

Source

Thrown at src/lightning/pytorch/loops/optimization/automatic.py:68

    loss: Optional[Tensor] = field(init=False, default=None)
    extra: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self._clone_loss()

    def _clone_loss(self) -> None:
        if self.closure_loss is not None:
            # the loss will get scaled for amp. avoid any modifications to it
            self.loss = self.closure_loss.detach().clone()

    @classmethod
    def from_training_step_output(cls, training_step_output: STEP_OUTPUT, normalize: int = 1) -> "ClosureResult":
        closure_loss, extra = None, {}

        if isinstance(training_step_output, Mapping):
            closure_loss = training_step_output.get("loss")
            if closure_loss is None:
                raise MisconfigurationException(
                    "In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present"
                )
            extra = {k: v for k, v in training_step_output.items() if k != "loss"}
        elif isinstance(training_step_output, Tensor):
            closure_loss = training_step_output
        elif training_step_output is not None:
            raise MisconfigurationException(
                "In automatic optimization, `training_step` must return a Tensor, a dict, or None (where the step will"
                " be skipped)."
            )

        if closure_loss is not None:
            # accumulate the loss. If ``accumulate_grad_batches == 1``, no effect
            # note: avoid in-place operation `x /= y` here on purpose
            closure_loss = closure_loss / normalize

        return cls(closure_loss, extra=extra)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Compute and include the loss: `return {'loss': loss, 'preds': logits}`
2. To skip a step legitimately, `return None` instead of a dict without loss
3. If you compute losses manually, set `self.automatic_optimization = False` in the model

Example fix

# before
def training_step(self, batch, batch_idx):
    logits = self(batch[0])
    return {'logits': logits}

# after
def training_step(self, batch, batch_idx):
    logits = self(batch[0])
    loss = self.loss_fn(logits, batch[1])
    return {'loss': loss, 'logits': logits}

out = self.training_step(batch, batch_idx)  # in a unit test
assert out is None or isinstance(out, torch.Tensor) or 'loss' in out

def has_loss_key(out) -> bool:
    return out is None or isinstance(out, torch.Tensor) or (isinstance(out, Mapping) and out.get('loss') is not None)

When it happens

Trigger: `def training_step(self, batch, batch_idx): return {'preds': logits, 'targets': y}` (no 'loss' key), or returning `{'loss': None}` conditionally; switching automatic_optimization=True (default) while the step was written for manual optimization.

Common situations: Returning metrics-only dicts from training_step; early-exit logic like `if batch is weird: return {'skipped': True}`; adapting a manual-optimization LightningModule to automatic without adding the loss key.

Related errors

• In automatic optimization, `training_step` must return a Ten
• Skipping backward by returning `None` from your `training_st
• `training_step` returned `None`. If this was on purpose, ign
• With `def training_step(self, dataloader_iter)`, `self.log(.
• Skipping the `training_step` by returning None in distribute

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/be3e965910628e39. Report an issue: GitHub.