{"record":{"id":"fd05dd0cb1dd63f2","repo":"Lightning-AI/pytorch-lightning","slug":"skipping-backward-by-returning-none-from-your-t-fd05dd","errorCode":null,"errorMessage":"Skipping backward by returning `None` from your `training_step` is not implemented with XLA. Please, open an issue in `https://github.com/Lightning-AI/pytorch-lightning/issues` requesting this feature.","messagePattern":"Skipping backward by returning `None` from your `training_step` is not implemented with XLA\\. Please, open an issue in `https://github\\.com/Lightning-AI/pytorch-lightning/issues` requesting this feature\\.","errorType":"exception","errorClass":"MisconfigurationException","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/plugins/precision/xla.py","lineNumber":80,"sourceCode":"    @override\n    def optimizer_step(  # type: ignore[override]\n        self,\n        optimizer: Optimizable,\n        model: \"pl.LightningModule\",\n        closure: Callable[[], Any],\n        **kwargs: Any,\n    ) -> Any:\n        import torch_xla.core.xla_model as xm\n\n        closure = partial(self._xla_wrap_closure, optimizer, closure)\n        closure = partial(self._wrap_closure, model, optimizer, closure)\n        closure_result = optimizer.step(closure=closure, **kwargs)\n        xm.mark_step()\n        skipped_backward = closure_result is None\n        # in manual optimization, the closure does not return a value\n        if model.automatic_optimization and skipped_backward:\n            # we lack coverage here so disable this - something to explore if there's demand\n            raise MisconfigurationException(\n                \"Skipping backward by returning `None` from your `training_step` is not implemented with XLA.\"\n                \" Please, open an issue in `https://github.com/Lightning-AI/pytorch-lightning/issues`\"\n                \" requesting this feature.\"\n            )\n        return closure_result\n\n    @override\n    def teardown(self) -> None:\n        os.environ.pop(\"XLA_USE_BF16\", None)\n        os.environ.pop(\"XLA_USE_F16\", None)\n\n    def _xla_wrap_closure(self, optimizer: Optimizable, closure: Callable[[], Any]) -> Any:\n        import torch_xla.core.xla_model as xm\n\n        closure_result = closure()\n        xm.reduce_gradients(optimizer)\n        return closure_result\n","sourceCodeStart":62,"sourceCodeEnd":98,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/plugins/precision/xla.py#L62-L98","documentation":"In automatic optimization, Lightning detects that training_step returned None (i.e. you skipped backward). The XLA precision plugin's optimizer_step explicitly rejects this because the XLA graph execution path (xm.mark_step) does not support skipping backward. The maintainers ask users to open a feature request if they need it.","triggerScenarios":"Your LightningModule.training_step returns None (or implicitly returns None) while automatic_optimization is True and you use the XLA precision plugin / XLA strategy; optimizer_step then raises MisconfigurationException.","commonSituations":"Porting code that conditionally skips batches (e.g. return None on empty batch) to TPU; refactoring training_step and accidentally dropping the loss return.","solutions":["Always return a loss tensor from training_step (never None) when using XLA with automatic optimization","If you must skip batches, gate the batch in train_dataloader/on_train_batch_start instead of returning None","If you genuinely need skipped backward on XLA, open the referenced GitHub issue requesting the feature"],"exampleFix":"# before\ndef training_step(self, batch, batch_idx):\n    if batch is None:\n        return None  # triggers MisconfigurationException on XLA\n    loss = self(batch).loss\n    self.log(\"loss\", loss)\n\n# after\ndef training_step(self, batch, batch_idx):\n    loss = self(batch).loss  # always compute and return loss\n    self.log(\"loss\", loss)\n    return loss","handlingStrategy":"validation","validationCode":"out = model.training_step(batch, batch_idx)\nassert out is not None, \"training_step must return a loss on XLA\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always return the loss from training_step in automatic optimization","Add a unit test asserting training_step returns a tensor","Filter problematic batches in dataloader or on_train_batch_start, not via None returns"],"tags":["xla","training-step","backward","pytorch-lightning","misconfiguration"],"backgroundTag":"unsupported-framework-operation","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}