{"record":{"id":"be3e965910628e39","repo":"Lightning-AI/pytorch-lightning","slug":"in-automatic-optimization-when-training-step-re","errorCode":null,"errorMessage":"In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present","messagePattern":"In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present","errorType":"exception","errorClass":"MisconfigurationException","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/loops/optimization/automatic.py","lineNumber":68,"sourceCode":"    loss: Optional[Tensor] = field(init=False, default=None)\n    extra: dict[str, Any] = field(default_factory=dict)\n\n    def __post_init__(self) -> None:\n        self._clone_loss()\n\n    def _clone_loss(self) -> None:\n        if self.closure_loss is not None:\n            # the loss will get scaled for amp. avoid any modifications to it\n            self.loss = self.closure_loss.detach().clone()\n\n    @classmethod\n    def from_training_step_output(cls, training_step_output: STEP_OUTPUT, normalize: int = 1) -> \"ClosureResult\":\n        closure_loss, extra = None, {}\n\n        if isinstance(training_step_output, Mapping):\n            closure_loss = training_step_output.get(\"loss\")\n            if closure_loss is None:\n                raise MisconfigurationException(\n                    \"In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present\"\n                )\n            extra = {k: v for k, v in training_step_output.items() if k != \"loss\"}\n        elif isinstance(training_step_output, Tensor):\n            closure_loss = training_step_output\n        elif training_step_output is not None:\n            raise MisconfigurationException(\n                \"In automatic optimization, `training_step` must return a Tensor, a dict, or None (where the step will\"\n                \" be skipped).\"\n            )\n\n        if closure_loss is not None:\n            # accumulate the loss. If ``accumulate_grad_batches == 1``, no effect\n            # note: avoid in-place operation `x /= y` here on purpose\n            closure_loss = closure_loss / normalize\n\n        return cls(closure_loss, extra=extra)\n","sourceCodeStart":50,"sourceCodeEnd":86,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/loops/optimization/automatic.py#L50-L86","documentation":"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.","triggerScenarios":"`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.","commonSituations":"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.","solutions":["Compute and include the loss: `return {'loss': loss, 'preds': logits}`","To skip a step legitimately, `return None` instead of a dict without loss","If you compute losses manually, set `self.automatic_optimization = False` in the model"],"exampleFix":"# before\ndef training_step(self, batch, batch_idx):\n    logits = self(batch[0])\n    return {'logits': logits}\n\n# after\ndef training_step(self, batch, batch_idx):\n    logits = self(batch[0])\n    loss = self.loss_fn(logits, batch[1])\n    return {'loss': loss, 'logits': logits}","handlingStrategy":"type-guard","validationCode":"out = self.training_step(batch, batch_idx)  # in a unit test\nassert out is None or isinstance(out, torch.Tensor) or 'loss' in out","typeGuard":"def has_loss_key(out) -> bool:\n    return out is None or isinstance(out, torch.Tensor) or (isinstance(out, Mapping) and out.get('loss') is not None)","tryCatchPattern":null,"preventionTips":["Always return {'loss': loss, ...} from training_step under automatic optimization","Add a fast dev run (Trainer(fast_dev_run=True)) in CI to exercise the step contract"],"tags":["pytorch-lightning","training-step","loss","automatic-optimization"],"backgroundTag":"missing-loss-key","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}