code: def test_checkpoint_wrapper_cpu_offload(self):
      model = nn.Sequential(
          nn.Linear(10, 10),
          nn.Linear(10, 10),
          nn.Linear(10, 10),
      ).cuda()
  
      # Patch saved_tensor_hooks to make the unpack keep the tensor on CPU for
      # testing, otherwise the tensor access during the DFS will cause orig
      # unpack to run, transferring the tensor back to GPU.
      def patched_init(saved_tensor_hook_obj, pack_hook, _):
          saved_tensor_hook_obj.pack_hook = pack_hook
  
          def testing_cpu_offload_unpack_hook(packed):
              _, tensor = packed
              return tensor
  
          saved_tensor_hook_obj.unpack_hook = (
              testing_cpu_offload_unpack_hook
          )
  
      orig_init = torch.autograd.graph.saved_tensors_hooks.__init__
      torch.autograd.graph.saved_tensors_hooks.__init__ = patched_init
  
      model = offload_wrapper(model)
  
      inp = torch.randn(3, 10, device="cuda")
      loss = model(inp).sum()
  
      # All autograd saved tensors should be offloaded to CPU.
      offload_verified = False
  
      def dfs(grad_fn):
          for e in dir(grad_fn):
              if not e.startswith(_SAVED_PREFIX):
                  continue
  
              saved = getattr(grad_fn, e)
              if isinstance(saved, torch.Tensor):
                  self.assertEqual(torch.device("cpu"), saved.device)
                  nonlocal offload_verified
                  offload_verified = True
  
          if hasattr(grad_fn, GRAD_FN_NEXT_FUNCTIONS):
              for next_grad_fn, _ in grad_fn.next_functions:
                  dfs(next_grad_fn)
  
      dfs(loss.grad_fn)
  
      self.assertTrue(offload_verified)
  
      torch.autograd.graph.saved_tensors_hooks.__init__ = orig_init
  
file_path: pytorch/test/distributed/fsdp/test_checkpoint_wrapper.py
imports: import contextlib
  import unittest
  from copy import deepcopy
  from functools import partial
  import torch
  import torch.nn
  from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import apply_activation_checkpointing, checkpoint_wrapper, CheckpointImpl, CheckpointWrapper, offload_wrapper, OffloadWrapper
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.testing._internal.common_utils import run_tests, TestCase
  from torch.utils.checkpoint import checkpoint
label_desc: Use image node_id 8 for calling the CheckpointWrapperTest obj's underlying member method code with example usage: obj.test_checkpoint_wrapper_cpu_offload() and returns: tensor