def test_deepspeed_configure_optimizer_device_set(tmp_path):
    """Test to ensure that the LM has access to the device within the ``configure_optimizer`` function, and
    estimated_stepping_batches works correctly as a result."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            assert self.trainer.estimated_stepping_batches == 1
            assert self.device.type == "cuda"
            raise SystemExit

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        fast_dev_run=True,
        accelerator="gpu",
        devices=2,
        strategy=DeepSpeedStrategy(),
    )
    with pytest.raises(SystemExit):
        trainer.fit(model)
Test to ensure that the LM has access to the device within the ``configure_optimizer`` function, and
    estimated_stepping_batches works correctly as a result.
tests/tests_pytorch/strategies/test_deepspeed.py