code: def checkpoint_wrapper(m, offload_to_cpu=False):
      """
      A friendlier wrapper for performing activation checkpointing.
  
      Compared to the PyTorch version, this version:
      - wraps an nn.Module, so that all subsequent calls will use checkpointing
      - handles keyword arguments in the forward
      - handles non-Tensor outputs from the forward
  
      Usage::
  
          checkpointed_module = checkpoint_wrapper(my_module, offload_to_cpu=True)
          a, b = checkpointed_module(x, y=3, z=torch.Tensor([1]))
      """
      # should I check whether original_forward has already been set?
      assert not hasattr(
          m, "precheckpoint_forward"
      ), "checkpoint function has already been applied?"
      m.precheckpoint_forward = m.forward
      m.forward = functools.partial(
          _checkpointed_forward,
          m.precheckpoint_forward,  # original_forward
          offload_to_cpu,
      )
      return m
  
file_path: fairseq/fairseq/modules/checkpoint_activations.py
imports: import functools
  from typing import Any, Dict, List, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from fairseq import utils
label_desc: Use image node_id 1 for calling a global function with example usage: checkpoint_wrapper(m, offload_to_cpu) and returns: m