Lightning-AI/pytorch-lightning · error · MisconfigurationException

In automatic optimization, `training_step` must return a Ten

Error message

In automatic optimization, `training_step` must return a Tensor, a dict, or None (where the step will be skipped).

What it means

Raised by ClosureResult.from_training_step_output in automatic optimization when training_step returns a value that is not a Tensor, not a Mapping, and not None — e.g. a tuple, list, number, or arbitrary object. The automatic optimization path can only derive a backward-able loss from those three shapes.

Source

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

        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)

    @override
    def asdict(self) -> dict[str, Any]:
        return {"loss": self.loss, **self.extra}


class Closure(AbstractClosure[ClosureResult]):
    """An implementation of a :class:`AbstractClosure` for automatic optimization in Lightning that combines three

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Return the loss tensor directly or a dict containing 'loss'
2. For extra outputs use `return {'loss': loss, 'logits': logits}`
3. Return None to intentionally skip the step

Example fix

# before
def training_step(self, batch, batch_idx):
    return loss.item(), logits  # tuple of float -> raises

# after
def training_step(self, batch, batch_idx):
    return {'loss': loss, 'logits': logits}

allowed = (torch.Tensor, Mapping, type(None))
assert isinstance(training_step_output, allowed) or training_step_output is None

def valid_step_output(out) -> bool:
    return out is None or isinstance(out, (torch.Tensor, Mapping))

When it happens

Trigger: `return loss.item(), logits` (tuple), `return [loss, logits]` (list), `return float(loss)`; returning a dataclass or namedtuple from training_step with automatic_optimization=True.

Common situations: Porting vanilla PyTorch training code that returns tuples; mixing up conventions with other frameworks (fastai, HF Trainer) where tuple returns are common; returning loss.item() which strips the graph.

Related errors

• In automatic_optimization, when `training_step` returns a di
• In manual optimization, `training_step` must either return a
• Skipping backward by returning `None` from your `training_st
• 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/8dbee8966f0cd035. Report an issue: GitHub.