def METHOD_NAME(self, function: Callable) -> Callable:
    # Working around a dispatcher issue with `torch.func.functionalize` when used
    # together with `make_fx`.
    # Ref: https://github.com/pytorch/pytorch/issues/99774#issuecomment-1527949391
    def wrapped(*inputs):
        inputs_functional = pytree.tree_map_only(
            torch.Tensor, torch._to_functional_tensor, inputs
        )
        torch._enable_functionalization(reapply_views=True)
        try:
            out = function(*inputs_functional)
        finally:
            torch._disable_functionalization()
        flat_inputs, _ = pytree.tree_flatten(inputs)
        flat_inputs_functional, _ = pytree.tree_flatten(inputs_functional)
        for inpt, input_functional in zip(flat_inputs, flat_inputs_functional):
            if isinstance(input_functional, torch.Tensor):
                torch._sync(input_functional)
                inpt_new = torch._from_functional_tensor(input_functional)
        pytree.tree_map(torch._sync, out)
        out_unwrapped = pytree.tree_map(torch._from_functional_tensor, out)
        return out_unwrapped
    return wrapped