def test_dataloader_reset_with_scale_batch_size(tmp_path, caplog, scale_method, init_batch_size):
    """Test that train and val dataloaders are reset at every update in scale batch size."""
    model = BatchSizeModel(batch_size=16)
    max_trials = 2
    scale_batch_size_kwargs = {
        "max_trials": max_trials,
        "steps_per_trial": 2,
        "init_val": init_batch_size,
        "mode": scale_method,
    }

    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1)
    tuner = Tuner(trainer)

    with caplog.at_level(logging.INFO):
        new_batch_size = tuner.scale_batch_size(model, **scale_batch_size_kwargs)

    dataset_len = len(trainer.train_dataloader.dataset)
    assert dataset_len == 64
    assert caplog.text.count("trying batch size") == (max_trials if init_batch_size < dataset_len else 0)
    assert caplog.text.count("greater or equal than the length") == int(new_batch_size == dataset_len)

    assert trainer.train_dataloader.batch_size == new_batch_size
    assert trainer.val_dataloaders.batch_size == new_batch_size
Test that train and val dataloaders are reset at every update in scale batch size.
tests/tests_pytorch/tuner/test_scale_batch_size.py