code: def test_checkpoint_wrapper_kwarg_support(self):
      class MyModel(nn.Module):
          def __init__(self):
              super().__init__()
              self.lin = nn.Linear(10, 10)
  
          def forward(self, a, b, c=None, d=None, **kwargs):
              return (
                  self.lin(a),
                  self.lin(b),
                  self.lin(c),
                  self.lin(d),
              )
  
      for wrapper in [
          partial(
              checkpoint_wrapper,
              checkpoint_impl=CheckpointImpl.REENTRANT,
          ),
          partial(
              checkpoint_wrapper,
              checkpoint_impl=CheckpointImpl.NO_REENTRANT,
          ),
          offload_wrapper,
      ]:
          with self.subTest(wrapper=wrapper):
              model = wrapper(MyModel())
              if wrapper == offload_wrapper:
                  self.assertTrue(isinstance(model, OffloadWrapper))
              else:
                  self.assertTrue(isinstance(model, CheckpointWrapper))
              # Verify kwargs can be passed in
              inp = torch.ones(4, 10, requires_grad=True)
              out = model(inp, inp, c=inp, d=inp, e=inp, f=inp)
              self.assertTrue(isinstance(out, tuple))
              self.assertEqual(4, len(out))
              # Without kwargs should have equivalent gradient requirements.
              out_no_kwarg = model(inp, inp, inp, inp)
              for t1, t2 in zip(out_no_kwarg, out):
                  self.assertEqual(t1, t2)
                  self.assertEqual(t1.requires_grad, t2.requires_grad)
  
      # Test model that enforces kwarg inputs
      class ModelEnforceKwarg(nn.Module):
          def __init__(self):
              super().__init__()
              self.lin = nn.Linear(10, 10)
  
          def forward(self, *, a=None, b=None):
              return (self.lin(a), self.lin(b))
  
      model = checkpoint_wrapper(
          ModelEnforceKwarg(), checkpoint_impl=CheckpointImpl.REENTRANT
      )
  
      inp = torch.ones(4, 10, requires_grad=True)
      out = model(a=inp, b=inp)
      self.assertEqual(2, len(out))
  
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 2 for calling the CheckpointWrapperTest obj's underlying member method code with example usage: obj.test_checkpoint_wrapper_kwarg_support() and returns: self,  self,  self,  self, self,  self