code: def test_checkpoint_requires_grad():
      """Test to check checkpointing when outputs do not require gradient."""
  
      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
              return self.linear(x)
  
      x = torch.rand(4, 2)
      model = nn.Sequential(
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
      )
      model[0].requires_grad_(False)
      model[1].requires_grad_(False)
      model[2].requires_grad_(False)
  
      y = model(x)
      y = y.sum()
      y.backward()
  
      # Since only last model needs grad, we only run forward twice for it
      assert model[0].cnt == 1
      assert model[1].cnt == 1
      assert model[2].cnt == 1
      assert model[3].cnt == 2
  
      # Now test with first model needing grad
      model = nn.Sequential(
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
          checkpoint_wrapper(TestModel()),
      )
      model[0].requires_grad_(True)
      model[1].requires_grad_(False)
      model[2].requires_grad_(False)
  
      y = model(x)
      y = y.sum()
      y.backward()
  
      # Since first model needs grad, all models need grad, so we run forward twice for all
      assert model[0].cnt == 2
      assert model[1].cnt == 2
      assert model[2].cnt == 2
      assert model[3].cnt == 2
  
      # Stress test with multiple inputs/outputs, of which some are not Tensor
      class TestModel2(nn.Module):
          def __init__(self):
              super().__init__()
              self.cnt = 0
              self.linear = nn.Linear(2, 2)
  
          def forward(self, x, y, z):
              self.cnt += 1
              z = z + [self.cnt]
              return self.linear(x + y), z, ["hi"]
  
      model1 = checkpoint_wrapper(TestModel())
      model2 = checkpoint_wrapper(TestModel())
      model3 = checkpoint_wrapper(TestModel2())
      model4 = checkpoint_wrapper(TestModel())
      model1.requires_grad_(False)
      model2.requires_grad_(False)
  
      y = model4(model3(model1(x), model2(x), ["bye"])[0])
      y = y.sum()
      y.backward()
  
      assert model1.cnt == 1
      assert model2.cnt == 1
      assert model3.cnt == 2
      assert model4.cnt == 2
  
      model1 = checkpoint_wrapper(TestModel())
      model2 = checkpoint_wrapper(TestModel())
      model3 = checkpoint_wrapper(TestModel2())
      model4 = checkpoint_wrapper(TestModel())
      model2.requires_grad_(False)
  
      y = model4(model3(model1(x), model2(x), ["bye"])[0])
      y = y.sum()
      y.backward()
  
      assert model1.cnt == 2
      assert model2.cnt == 1
      assert model3.cnt == 2
      assert model4.cnt == 2
  
      # Test flattened pararameters
      model = nn.Sequential(
          FlattenParamsWrapper(checkpoint_wrapper(TestModel())),
          FlattenParamsWrapper(checkpoint_wrapper(TestModel())),
          FlattenParamsWrapper(checkpoint_wrapper(TestModel())),
          FlattenParamsWrapper(checkpoint_wrapper(TestModel())),
      )
      model[0].requires_grad_(False)
      model[1].requires_grad_(False)
  
      y = model(x)
      y = y.sum()
      y.backward()
  
      assert model[0].cnt == 1
      assert model[1].cnt == 1
      assert model[2].cnt == 2
      assert model[3].cnt == 2
  
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 9 for calling a global function with example usage: test_checkpoint_requires_grad() and returns: self, self,  z,  list