code: def test_load_activation_checkpointed_module(self):
      lin = nn.Linear(10, 10, bias=False)
      lin = checkpoint_wrapper(
          lin,
          checkpoint_fn=checkpoint,
          # checkpoint kwargs
          use_reentrant=True,
          preserve_rng_state=False,
      )
      state_dict = deepcopy(lin.state_dict())
      # Load into non-checkpoint wrapped linear module
      lin_new = nn.Linear(10, 10, bias=False)
      lin_new.load_state_dict(state_dict)
      for p1, p2 in zip(lin.parameters(), lin_new.parameters()):
          self.assertEqual(p1, p2)
          self.assertTrue(torch.allclose(p1, p2))
  
      # Load non-checkpoint wrapped module into checkpoint wrapped one
      # Make params different
      for p in lin_new.parameters():
          with torch.no_grad():
              p.add_(0.5)
  
      state_dict = deepcopy(lin_new.state_dict())
      # Verify checkpoint wrapped linear can load unwrapped linear
      lin.load_state_dict(state_dict)
      for p1, p2 in zip(lin.parameters(), lin_new.parameters()):
          self.assertEqual(p1, p2)
  
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 1 for calling the CheckpointWrapperTest obj's underlying member method code with example usage: obj.test_load_activation_checkpointed_module() without return types