{"record":{"id":"8dbee8966f0cd035","repo":"Lightning-AI/pytorch-lightning","slug":"in-automatic-optimization-training-step-must-re","errorCode":null,"errorMessage":"In automatic optimization, `training_step` must return a Tensor, a dict, or None (where the step will be skipped).","messagePattern":"In automatic optimization, `training_step` must return a Tensor, a dict, or None \\(where the step will be skipped\\)\\.","errorType":"exception","errorClass":"MisconfigurationException","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/loops/optimization/automatic.py","lineNumber":75,"sourceCode":"        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\n    @override\n    def asdict(self) -> dict[str, Any]:\n        return {\"loss\": self.loss, **self.extra}\n\n\nclass Closure(AbstractClosure[ClosureResult]):\n    \"\"\"An implementation of a :class:`AbstractClosure` for automatic optimization in Lightning that combines three","sourceCodeStart":57,"sourceCodeEnd":93,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/loops/optimization/automatic.py#L57-L93","documentation":"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.","triggerScenarios":"`return loss.item(), logits` (tuple), `return [loss, logits]` (list), `return float(loss)`; returning a dataclass or namedtuple from training_step with automatic_optimization=True.","commonSituations":"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.","solutions":["Return the loss tensor directly or a dict containing 'loss'","For extra outputs use `return {'loss': loss, 'logits': logits}`","Return None to intentionally skip the step"],"exampleFix":"# before\ndef training_step(self, batch, batch_idx):\n    return loss.item(), logits  # tuple of float -> raises\n\n# after\ndef training_step(self, batch, batch_idx):\n    return {'loss': loss, 'logits': logits}","handlingStrategy":"type-guard","validationCode":"allowed = (torch.Tensor, Mapping, type(None))\nassert isinstance(training_step_output, allowed) or training_step_output is None","typeGuard":"def valid_step_output(out) -> bool:\n    return out is None or isinstance(out, (torch.Tensor, Mapping))","tryCatchPattern":null,"preventionTips":["Standardize on returning a dict with 'loss' for automatic optimization","Write a unit test asserting the return type of training_step"],"tags":["pytorch-lightning","training-step","return-type","automatic-optimization"],"backgroundTag":"invalid-training-step-return-type","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}