def test_iterable_dataset_stop_iteration_at_epoch_beginning(yield_at_all, tmp_path):
    """Test that the training loop skips execution if the iterator is empty from the start."""

    class TestDataset(IterableDataset):
        def __init__(self, gen):
            self.gen = gen

        def __iter__(self):
            return iter(self.gen())

    class TestModel(BoringModel):
        def gen(self):
            # produce data in epoch 0, no data otherwise
            if yield_at_all and self.current_epoch == 0:
                yield torch.rand(32)
                yield torch.rand(32)
                yield torch.rand(32)

    model = TestModel()
    train_dataloader = DataLoader(TestDataset(model.gen), batch_size=2)
    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=False,
        max_epochs=2,
        enable_model_summary=False,
    )
    trainer.fit(model, train_dataloaders=train_dataloader)
    assert trainer.global_step == 2 * yield_at_all
    # even though the generator might not yield any data, the fit_loop still advances so the
    # current epoch gets increased
    assert trainer.current_epoch == 2
Test that the training loop skips execution if the iterator is empty from the start.
tests/tests_pytorch/trainer/test_dataloaders.py