def test_reapply_compile():
    """Test that Fabric can rewrap a compiled module such that compilation happens over the DDP-wrapper."""
    fabric = Fabric(accelerator="cuda", devices=2, strategy="ddp")
    fabric.launch()

    model = BoringModel()
    # currently (PyTorch 2.6) using ruduce-overhead here casues a RuntimeError:
    # Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.
    compile_kwargs = {"mode": "reduce-overhead"} if _TORCH_LESS_EQUAL_2_6 else {}
    compiled_model = torch.compile(model, **compile_kwargs)
    torch.compile.reset_mock()

    fabric_model = fabric.setup(compiled_model, _reapply_compile=True)

    assert isinstance(fabric_model._forward_module, OptimizedModule)
    assert isinstance(fabric_model._forward_module._orig_mod, DistributedDataParallel)
    # Assert we called compile again with the same arguments, but on the DDP-wrapped module
    torch.compile.assert_called_with(fabric_model._forward_module._orig_mod, **compile_kwargs)

    assert fabric_model._original_module == model
    assert fabric_model._forward_module._orig_mod.module == model
    assert fabric_model.device == fabric.device

    # Smoke-testing forward to ensure we don't get compilation errors
    for _ in range(3):
        fabric_model(torch.randn(2, 32, device=fabric.device)).sum().backward()
Test that Fabric can rewrap a compiled module such that compilation happens over the DDP-wrapper.
tests/tests_fabric/strategies/test_ddp_integration.py