def test_setup_compiled_module(reapply_compile, setup_method):
    """Test that an `OptimizedModule` can be passed to the setup method."""
    from torch._dynamo.eval_frame import OptimizedModule

    fabric = Fabric(devices=1)
    model = nn.Linear(1, 2)
    compiled_model = torch.compile(model)
    assert compiled_model._compile_kwargs is not None
    assert isinstance(compiled_model, OptimizedModule)
    setup_method = getattr(fabric, setup_method)
    fabric_model = setup_method(compiled_model, _reapply_compile=reapply_compile)

    assert isinstance(fabric_model._forward_module, OptimizedModule)
    if reapply_compile:
        # The forward_module got rewrapped into a new OptimizedModule
        assert fabric_model._forward_module != fabric_model._original_module
        # The original_module points to the pure module
        assert fabric_model._original_module is model
        assert fabric_model._forward_module._orig_mod is model
    else:
        assert fabric_model._forward_module is fabric_model._original_module
    # Attributes get passed through
    assert fabric_model.weight is model.weight
Test that an `OptimizedModule` can be passed to the setup method.
tests/tests_fabric/test_fabric.py