code: def _test_checkpoint_wrapper(self, device, log_memory_usage=False):
      def get_loss_and_gnorm(model):
          torch.manual_seed(1)
          input = torch.rand(2, 16, 32).requires_grad_(True).to(device)
          model.zero_grad()
          loss = model(input).sum()
          loss.backward()
          gnorm = torch.norm(
              torch.stack(
                  [
                      torch.norm(p.grad.detach())
                      for p in model.parameters()
                  ]
              )
          )
          return {"loss": loss, "gnorm": gnorm}
  
      model = Model().to(device)
      no_cpt = get_loss_and_gnorm(model)
  
      model = Model(use_pytorch_checkpoint=True).to(device)
      pyt_cpt = get_loss_and_gnorm(model)
      torch.testing.assert_allclose(no_cpt["loss"], pyt_cpt["loss"])
      torch.testing.assert_allclose(no_cpt["gnorm"], pyt_cpt["gnorm"])
  
      model = Model(use_fairseq_checkpoint=True).to(device)
      fairseq_cpt = get_loss_and_gnorm(model)
      torch.testing.assert_allclose(no_cpt["loss"], fairseq_cpt["loss"])
      torch.testing.assert_allclose(
          no_cpt["gnorm"], fairseq_cpt["gnorm"]
      )
  
      model = Model(
          use_fairseq_checkpoint=True, offload_to_cpu=True
      ).to(device)
      fairseq_cpt_offload = get_loss_and_gnorm(model)
      torch.testing.assert_allclose(
          no_cpt["loss"], fairseq_cpt_offload["loss"]
      )
      torch.testing.assert_allclose(
          no_cpt["gnorm"], fairseq_cpt_offload["gnorm"]
      )
  
file_path: fairseq/tests/test_activation_checkpointing.py
imports: import unittest
  import torch
  import torch.nn
  from fairseq.modules.checkpoint_activations import checkpoint_wrapper
  from torch.utils.checkpoint import checkpoint
label_desc: Use image node_id 1 for calling the TestActivationCheckpointing obj's underlying member method code with example usage: obj._test_checkpoint_wrapper(device, log_memory_usage) and returns: dict