def test_num_stepping_batches_gpu(trainer_kwargs, estimated_steps, monkeypatch):
    """Test stepping batches with GPU strategies."""
    num_devices_per_node = 7
    mock_cuda_count(monkeypatch, num_devices_per_node)
    trainer = Trainer(max_epochs=1, devices=num_devices_per_node, accelerator="gpu", **trainer_kwargs)

    # set the `parallel_devices` to cpu to run the test on CPU and take `num_nodes`` into consideration
    # because we can't run on multi-node in ci
    trainer.strategy.parallel_devices = [torch.device("cpu", index=i) for i in range(num_devices_per_node)]

    model = BoringModel()
    trainer._data_connector.attach_data(model)
    trainer.strategy.connect(model)
    assert trainer.estimated_stepping_batches == estimated_steps
Test stepping batches with GPU strategies.
tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py