code: def _checkpointed_forward(
      original_forward: Any,
      weak_self: Any,
      offload_to_cpu: bool,
      *args: Any,
      **kwargs: Any,
  ) -> Any:
      module = weak_self()
  
      # If gradients are disabled, just use original `.forward()` method directly.
      if (
          not torch.is_grad_enabled()
          or thread_local.is_checkpointing_disabled
      ):
          return original_forward(module, *args, **kwargs)
  
      # Autograd Functions in PyTorch work best with positional args, since
      # the backward must return gradients (or None) for every input argument.
      # We can flatten keyword arguments to make this easier.
      args = (module,) + args
      kwarg_keys, flat_args = pack_kwargs(*args, **kwargs)
      parent_ctx_dict: Dict[str, Any] = {
          "offload": offload_to_cpu,
      }
      # Dummy tensor with grad is used to ensure the backward pass is called. This is needed
      # when original_forward's input are non-tensor (i.e. a tuple). Using this dummy tensor
      # avoids requiring users to set their input tensors's requires_grad flag. In the case
      # of tuple type inputs, setting the flag won't even trigger the backward pass.
      #
      # One implication of this is that since we always feed in a dummy tensor
      # needing grad, then the output will always require grad, even if it originally
      # wouldn't, such as if the module and original input both do not require grad.
      # We get around this by saving the desired requires_grad value in output and
      # detaching the output if needed.
      output = CheckpointFunction.apply(
          torch.tensor([], requires_grad=True),
          original_forward,
          parent_ctx_dict,
          kwarg_keys,
          *flat_args,
      )
      output_requires_grad = parent_ctx_dict["output_requires_grad"]
      if not isinstance(output, torch.Tensor):
          # If output should not require grad, then detach it, since otherwise it will
          # always have requires_grad = True due to our dummy tensor input above that
          # requires_grad
          output = [
              x.detach() if not output_requires_grad else x
              for x in output
          ]
  
          packed_non_tensor_outputs = parent_ctx_dict[
              "packed_non_tensor_outputs"
          ]
          if packed_non_tensor_outputs:
              output = unpack_non_tensors(
                  output, packed_non_tensor_outputs
              )
  
      else:
          # If output should not require grad, then detach it, since otherwise it will
          # always have requires_grad = True due to our dummy tensor input above that
          # requires_grad
          if not output_requires_grad:
              output = output.detach()
  
      return output
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  import functools
  import threading
  from typing import Any, Dict, Generator, Optional, Tuple
  import weakref
  import torch
  from torch import Tensor
  import torch.nn
  import torch.utils.checkpoint
  from fairscale.internal.containers import pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
  from .checkpoint_utils import patch_batchnorm
label_desc: Use image node_id 7 for calling a global function with example usage: _checkpointed_forward(original_forward, weak_offload_to_cpu) and returns: output, original_forward