def test_training_step_no_return_when_even(tmp_path):
    """Tests correctness when some training steps have been skipped."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            self.training_step_called = True
            loss = self.step(batch[0])
            self.log("a", loss, on_step=True, on_epoch=True)
            return loss if batch_idx % 2 else None

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=4,
        limit_val_batches=1,
        max_epochs=4,
        enable_model_summary=False,
        logger=False,
        enable_checkpointing=False,
    )

    Closure.warning_cache.clear()

    with pytest.warns(UserWarning, match=r".*training_step` returned `None.*"):
        trainer.fit(model)

    trainer.state.stage = RunningStage.TRAINING

    # manually check a few batches
    for batch_idx, batch in enumerate(model.train_dataloader()):
        kwargs = {"batch": batch, "batch_idx": batch_idx}
        out = trainer.fit_loop.epoch_loop.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
        if not batch_idx % 2:
            assert out == {}
Tests correctness when some training steps have been skipped.
tests/tests_pytorch/loops/test_training_loop_flow_scalar.py