def METHOD_NAME(model, checkpoint_wrapper_fn=CheckpointWrapper, check_fn=lambda _: True):
    from torch.distributed.fsdp.wrap import _recursive_wrap
    _recursive_wrap(
        module=model,
        auto_wrap_policy=partial(lambda_auto_wrap_policy, lambda_fn=check_fn),
        wrapper_cls=checkpoint_wrapper_fn,
        ignored_modules=set(),
        ignored_params=set(),
        only_wrap_children=True,
    )