Lightning-AI/pytorch-lightning · warning

`training_step` returned `None`. If this was on purpose, ign

Error message

`training_step` returned `None`. If this was on purpose, ignore this warning...

What it means

AutomaticOptimization's closure warns when training_step returns None (closure_loss is None), meaning no loss was produced so no backward/step occurs for that batch. It is a warning only if intentional (e.g. skipping batches).

Source

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

    def __init__(
        self,
        step_fn: Callable[[], ClosureResult],
        backward_fn: Optional[Callable[[Tensor], None]] = None,
        zero_grad_fn: Optional[Callable[[], None]] = None,
    ):
        super().__init__()
        self._step_fn = step_fn
        self._backward_fn = backward_fn
        self._zero_grad_fn = zero_grad_fn

    @override
    @torch.enable_grad()
    def closure(self, *args: Any, **kwargs: Any) -> ClosureResult:
        step_output = self._step_fn()

        if step_output.closure_loss is None:
            self.warning_cache.warn("`training_step` returned `None`. If this was on purpose, ignore this warning...")

        if self._zero_grad_fn is not None:
            self._zero_grad_fn()

        if self._backward_fn is not None and step_output.closure_loss is not None:
            self._backward_fn(step_output.closure_loss)

        return step_output

    @override
    def __call__(self, *args: Any, **kwargs: Any) -> Optional[Tensor]:
        self._result = self.closure(*args, **kwargs)
        return self._result.loss


_OUTPUTS_TYPE = dict[str, Any]

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Ensure training_step returns the loss (or dict with 'loss' key) on the training path
  2. If skipping is intentional, keep the warning or return a zero-attached loss pattern
  3. Inspect step_output.closure_loss in on_train_batch_end when debugging

Example fix

# before
def training_step(self, batch, i):
    loss = self.loss(batch)
    # forgot return
# after
def training_step(self, batch, i):
    loss = self.loss(batch)
    return loss
Defensive patterns

Strategy: validation

Validate before calling

out = model.training_step(batch, 0)
assert out is not None and (torch.is_tensor(out) or 'loss' in out), 'training_step must return loss'

Type guard

def returns_loss(step_out) -> bool:
    import torch
    return torch.is_tensor(step_out) or (isinstance(step_out, dict) and step_out.get('loss') is not None)

Prevention

When it happens

Trigger: training_step with conditional return (returning None on skip), a bug like forgetting return loss, or returning a dict/tensor shape the loss extraction can't read ('loss' key missing yields None closure loss path).

Common situations: Gradient-accumulation or dynamic batch skipping logic; refactoring training_step and dropping the return statement.

Related errors


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