code: def _checkpointed_forward(
      original_forward, offload_to_cpu, *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.
      kwarg_keys, flat_args = pack_kwargs(*args, **kwargs)
      parent_ctx_dict = {"offload": offload_to_cpu}
      output = CheckpointFunction.apply(
          original_forward, parent_ctx_dict, kwarg_keys, *flat_args
      )
      if isinstance(output, torch.Tensor):
          return output
      else:
          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
              )
          return output
  
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 3 for calling a global function with example usage: _checkpointed_forward(original_forward, offload_to_cpu) and returns: output, output