docstr: Tests that checkpoint_wrapper can pass down args / kwargs to configure
  torch.utils.checkpoint.
docstr_tok: ["Tests","that","checkpoint_wrapper","can","pass","down","args","\/","kwargs","to","configure","torch.utils.checkpoint","."]
code: def test_checkpoint_wrapper_args_kwargs(self):
      """
      Tests that checkpoint_wrapper can pass down args / kwargs to configure
      torch.utils.checkpoint.
      """
  
      count = 0
  
      @contextlib.contextmanager
      def ctx_manager():
          nonlocal count
          count += 1
          yield
  
      def get_ctx_mgrs():
          return (ctx_manager(), ctx_manager())
  
      # kwargs test
      torch_utils_checkpoint = torch.utils.checkpoint.checkpoint
      m = checkpoint_wrapper(
          torch.nn.Linear(1, 1),
          checkpoint_fn=torch_utils_checkpoint,
          use_reentrant=False,
          context_fn=get_ctx_mgrs,
      )
      m(torch.randn(2, 1)).sum().backward()
      self.assertEqual(2, count)
  
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 3 for calling the CheckpointWrapperTest obj's underlying member method code with example usage: obj.test_checkpoint_wrapper_args_kwargs() and returns: ctx_manager,  ctx_manager