def test_strategy_checkpoint(state_dict_type, precision, tmp_path):
    """Test to ensure that checkpoint is saved correctly when using a single GPU, and all stages can be run."""
    model = TestFSDPModel()
    strategy = FSDPStrategy(state_dict_type=state_dict_type)
    trainer = Trainer(
        default_root_dir=tmp_path, accelerator="gpu", devices=2, strategy=strategy, precision=precision, max_epochs=1
    )
    _run_multiple_stages(trainer, model, os.path.join(tmp_path, "last.ckpt"))
Test to ensure that checkpoint is saved correctly when using a single GPU, and all stages can be run.
tests/tests_pytorch/strategies/test_fsdp.py