def test_deepspeed_multigpu_test_rnn(tmp_path):
    """Test to ensure that turning off explicit partitioning of the entire module for ZeRO Stage 3 works when training
    with certain layers which will crash with explicit partitioning."""

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.rnn = torch.nn.GRU(32, 32)

        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 turning off explicit partitioning of the entire module for ZeRO Stage 3 works when training
    with certain layers which will crash with explicit partitioning.
tests/tests_pytorch/strategies/test_deepspeed.py