def _capture_compile_kwargs(compile_fn: Callable) -> Callable:
    """Wraps the ``torch.compile`` function and captures the compile arguments.

    We extract the compile arguments so that we can reapply ``torch.compile`` in ``Fabric.setup()`` with the
    same arguments as the user passed to the original call. The arguments get stored in a dictionary
    ``_compile_kwargs`` on the returned compiled module.

    """
    # Limitation: Currently, the global compile config does not get captured on a per-model basis.
    # PyTorch will resolve this in the future: https://github.com/pytorch/pytorch/issues/116575

    @wraps(compile_fn)
    def _capture(*args: Any, **kwargs: Any) -> Any:
        if not args or not isinstance(args[0], nn.Module):
            # either torch.compile is being applied as a decorator or we're compiling something else
            return compile_fn(*args, **kwargs)

        model = args[0]
        compiled_model = compile_fn(model, **kwargs)
        compiled_model._compile_kwargs = deepcopy(kwargs)
        return compiled_model

    return _capture
Wraps the ``torch.compile`` function and captures the compile arguments.

    We extract the compile arguments so that we can reapply ``torch.compile`` in ``Fabric.setup()`` with the
    same arguments as the user passed to the original call. The arguments get stored in a dictionary
    ``_compile_kwargs`` on the returned compiled module.

    
src/lightning/fabric/wrappers.py