class_docstr: Helper class to wrap modules based on default config args via a context manager.
  See :func:`enable_wrap` for more information.
class_docstr_tok: ["Helper","class","to","wrap","modules","based","on","default","config","args","via","a","context","manager",".","See",":","func",":","`","enable_wrap","`","for","more","information","."]
code: def enable_autowrap_context(
      auto_wrap_policy: Optional[Callable], kwargs: Any
  ) -> None:
      if ConfigAutoWrap.in_autowrap_context:
          raise NotImplementedError(
              "You are already within an autowrap context and we currently do not supported nested autowrap."
          )
      ConfigAutoWrap.in_autowrap_context = True
      # Get and save the wrapper cls for the context.
      if "move_module_cuda_half" in kwargs.keys():
          ConfigAutoWrap.move_module_cuda_half = cast(
              bool, kwargs["move_module_cuda_half"]
          )
          del kwargs["move_module_cuda_half"]
      assert "wrapper_cls" in kwargs.keys()
      ConfigAutoWrap.wrapper_cls = cast(Callable, kwargs["wrapper_cls"])
      del kwargs["wrapper_cls"]
      # Save the rest.
      ConfigAutoWrap.auto_wrap_policy = (
          default_auto_wrap_policy
          if auto_wrap_policy is None
          else auto_wrap_policy
      )
      ConfigAutoWrap.kwargs = kwargs
  
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 2 for calling the ConfigAutoWrap obj's underlying member method code with example usage: obj.enable_autowrap_context(auto_wrap_policy, kwargs) without return types