def test_fsdp_v1_modules_unsupported():
    """Test that the strategy won't allow setting up a module wrapped with the legacy FSDP API."""
    from torch.distributed.fsdp import FullyShardedDataParallel

    module = Mock(modules=Mock(return_value=[Mock(spec=FullyShardedDataParallel)]))
    strategy = ModelParallelStrategy(parallelize_fn=(lambda x, _: x))
    with pytest.raises(TypeError, match="only supports the new FSDP2 APIs in PyTorch >= 2.4"):
        strategy.setup_module(module)
Test that the strategy won't allow setting up a module wrapped with the legacy FSDP API.
tests/tests_fabric/strategies/test_model_parallel.py