code: def enable_wrap(
      auto_wrap_policy: Optional[Callable] = None, **wrapper_kwargs: Any
  ) -> Generator[None, None, None]:
      """
      Context manager to wrap modules using a wrapper.
  
      Useful for when you'd like to apply the same parameters to all child modules
      that you wrap. A particularly important use case is wrapping large layers so
      that they get sharded (in-place) during initialization, to avoid running out of
      system memory. Large layers can indicate that they should be sharded via
      the ``wrap`` annotation and this context manager can provide the
      exact configuration for these nested instances.
  
      Usage::
  
          with enable_wrap(**params):
              # Wraps layer in FSDP by default if within context
              self.l1 = wrap(torch.nn.Linear(5, 5))
              self.l2 = auto_wrap(
                  TransformerBlock(),
                  # Wraps children modules based on a different min_num_params
                  auto_wrap_policy=functools.partial(default_auto_wrap_policy, min_num_params=1e7)
              )
  
      Args:
          auto_wrap_policy (Callable, Optional):
              Custom function to control how to do :func:`auto_wrap`. This is
              useful to exclude unsupported modules or wrap based on sizes when
              wrapping recursively. Note: modules annotated with :func:`wrap`
              ignore this policy and will always be wrapped.
              (default: :func:`default_auto_wrap_policy`)
          **wrapper_kwargs:
              Configuration settings that will be passed to all ``wrap``
              instances inside the context
      """
      with ConfigAutoWrap(auto_wrap_policy, **wrapper_kwargs):
          yield
  
file_path: fairscale/fairscale/nn/wrap/auto_wrap.py
imports: import contextlib
  from typing import Any, Callable, Dict, Generator, Optional, Set, Tuple, Type, cast
  import torch.nn
label_desc: Use image node_id 3 for calling a global function with example usage: enable_wrap(auto_wrap_policy) without return types