docstr: Wrap a module so that the forward method uses PyTorch's [checkpointing functionality]
  (https://pytorch.org/docs/stable/checkpoint.html).
  
  !!! Note
      Currently this `CheckpointWrapper` implementation requires that the wrapped
      module is called with positional arguments only.
  
      We recommend you use the
      :class:`allennlp.nn.checkpoint.fairscale_checkpoint_wrapper.FairScaleCheckpointWrapper`
      if you need more flexibility.
docstr_tok: ["Wrap","a","module","so","that","the","forward","method","uses","PyTorch","'s","[","checkpointing","functionality","]","(","https",":","\/\/pytorch.org\/docs\/stable\/checkpoint.html",")",".","!","!","!","Note","Currently","this","`","CheckpointWrapper","`","implementation","requires","that","the","wrapped","module","is","called","with","positional","arguments","only",".","We","recommend","you","use","the",":","class",":","`","allennlp.nn.checkpoint.fairscale_checkpoint_wrapper.FairScaleCheckpointWrapper","`","if","you","need","more","flexibility","."]
code: def wrap_module(self, module: nn.Module, **kwargs) -> nn.Module:
      """
      Wrap a module so that the forward method uses PyTorch's [checkpointing functionality]
      (https://pytorch.org/docs/stable/checkpoint.html).
  
      !!! Note
          Currently this `CheckpointWrapper` implementation requires that the wrapped
          module is called with positional arguments only.
  
          We recommend you use the
          :class:`allennlp.nn.checkpoint.fairscale_checkpoint_wrapper.FairScaleCheckpointWrapper`
          if you need more flexibility.
  
      """
      # Inspired by FairScale:
      #  --> https://github.com/facebookresearch/fairscale/blob/1e4a503cda8571851a68effd6e504a192838ab06/fairscale/nn/checkpoint/checkpoint_activations.py#L145-L153  # noqa: E501
      # We just patch the forward method to avoid having to proxy all the fields and other methods.
      # The use of weakref here is to prevent creating a ref cycle: m -> m.forward -> m.
  
      assert (
          len(kwargs) == 0
      )  # This way of wrapping only works for positional arguments.
  
      module.forward = functools.partial(  # type: ignore[assignment]
          _checkpointed_forward,
          type(module).forward,
          weakref.ref(module),
      )
      return module
  
file_path: allennlp/allennlp/nn/checkpoint/checkpoint_wrapper.py
imports: import functools
  from typing import Callable, TypeVar
  import weakref
  import torch
  import torch.nn
  from torch.utils.checkpoint import CheckpointFunction
  from allennlp.common.registrable import Registrable
label_desc: Use image node_id 1 for calling the TorchCheckpointWrapper obj's underlying member method code with example usage: obj.wrap_module(module) and returns: module