"module" should be wrapped if "recurse=False" or if the
        traversal should continue down the subtree if "recurse=True".
        Additional custom arguments may be added to the callable. The
        "size_based_auto_wrap_policy" in
        "torch.distributed.fsdp.wrap.py" gives an example callable
        that wraps a module if the parameters in its subtree exceed
        100M numel. A good practice is to print the model after
        wrapping and adjust as needed.
    Example:

       >>> def custom_auto_wrap_policy(
       >>>     module: nn.Module,
       >>>     recurse: bool,
       >>>     nonwrapped_numel: int,
       >>>     # Additional custom arguments
       >>>     min_num_params: int = int(1e8),
       >>> ) -> bool:
       >>>     return nonwrapped_numel >= min_num_params
       >>> # Configure a custom `min_num_params`
       >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
