def test_training_step_none_batches(tmp_path):
    """Tests correctness when the train dataloader gives None for some steps."""

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.counter = 0

        def collate_none_when_even(self, batch):
            result = None if self.counter % 2 == 0 else default_collate(batch)
            self.counter += 1
            return result

        def train_dataloader(self):
            return DataLoader(RandomDataset(32, 4), collate_fn=self.collate_none_when_even)

        def on_train_batch_end(self, outputs, batch, batch_idx):
            if batch_idx % 2 == 0:
                assert outputs is None
            else:
                assert outputs

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

    with pytest.warns(UserWarning, match=r".*train_dataloader yielded None.*"):
        trainer.fit(model)
Tests correctness when the train dataloader gives None for some steps.
tests/tests_pytorch/loops/test_training_loop_flow_scalar.py