def test_deepspeed_multigpu_partial_partition_parameters(tmp_path):
    """Test to ensure that a module that defines a layer inside the ``__init__`` and ``configure_model`` correctly
    converts all parameters to float16 when ``precision=16`` and runs successfully."""

    class TestModel(ModelParallelBoringModel):
        def __init__(self):
            super().__init__()
            self.layer_2 = torch.nn.Linear(32, 32)

        def configure_model(self) -> None:
            if self.layer is None:
                self.layer = torch.nn.Linear(32, 2)

        def forward(self, x):
            x = self.layer_2(x)
            return self.layer(x)

        def on_train_epoch_start(self) -> None:
            assert all(x.dtype == torch.float16 for x in self.parameters())

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        strategy=DeepSpeedStrategy(stage=3),
        accelerator="gpu",
        devices=1,
        fast_dev_run=True,
        precision="16-mixed",
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.fit(model)
Test to ensure that a module that defines a layer inside the ``__init__`` and ``configure_model`` correctly
    converts all parameters to float16 when ``precision=16`` and runs successfully.
tests/tests_pytorch/strategies/test_deepspeed.py