def test_strategy_sync_batchnorm(tmp_path):
    """Test to ensure that sync_batchnorm works when using FSDP and GPU, and all stages can be run."""
    model = TestFSDPModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="gpu",
        devices=2,
        strategy="fsdp",
        precision="16-mixed",
        max_epochs=1,
        sync_batchnorm=True,
    )
    _run_multiple_stages(trainer, model, os.path.join(tmp_path, "last.ckpt"))
Test to ensure that sync_batchnorm works when using FSDP and GPU, and all stages can be run.
tests/tests_pytorch/strategies/test_fsdp.py