code: def test_checkpoint_disabling():
      """Test to check new disable_checkpoint() API added to checkpoint_wrapper."""
  
      class TestModel(nn.Module):
          def __init__(self):
              super().__init__()
              self.cnt = 0
              self.linear = nn.Linear(2, 2)
  
          def forward(self, x):
              self.cnt += 1
              y = []
              for i in x:
                  y.append(self.linear(i))
              return y
  
      x = torch.rand(4, 2)
      model1 = checkpoint_wrapper(TestModel())
      model2 = checkpoint_wrapper(TestModel())
  
      # Forward. cnt += 1
      y = model1(x)
      y = sum(i.sum() for i in y)
      # Backward. cnt += 1
      y.backward()
      assert model1.cnt == 2
  
      with disable_checkpointing():
          # Forward. cnt += 1
          y = model2(x)
          y = sum(i.sum() for i in y)
          # Backward. cnt remains same as checkpointing is disabled
          y.backward()
      assert model2.cnt == 1
  
imports: import pytest
  import torch
  import torch.nn
  from torch.utils.checkpoint import checkpoint
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper, disable_checkpointing
  from fairscale.nn.misc import FlattenParamsWrapper
  from fairscale.nn.misc import checkpoint_wrapper
label_desc: Use image node_id 8 for calling a global function with example usage: test_checkpoint_disabling() and returns: y