def test_tuner_with_evaluation_methods(tmp_path, trainer_fn):
    """Test batch size tuner with Trainer's evaluation methods."""
    before_batch_size = 2
    max_trials = 4
    expected_scaled_batch_size = before_batch_size ** (max_trials + 1)

    model = BatchSizeModel(batch_size=before_batch_size)
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=100)
    tuner = Tuner(trainer)
    tuner.scale_batch_size(model, max_trials=max_trials, batch_arg_name="batch_size", method=trainer_fn)

    after_batch_size = model.batch_size
    loop = getattr(trainer, f"{trainer_fn}_loop")

    assert trainer.global_step == 0
    assert trainer.current_epoch == 0
    assert loop.batch_progress.current.completed == 0
    assert expected_scaled_batch_size == after_batch_size
    assert not any(f for f in os.listdir(tmp_path) if f.startswith(".scale_batch_size_temp_model"))
Test batch size tuner with Trainer's evaluation methods.
tests/tests_pytorch/tuner/test_scale_batch_size.py