code: def test_state_checkpointing():
      """Test state checkpointing on GPU since that's the common case.
  
      Note, we don't support checkpointing in the middle of gradient accumulation
      step. Therefore, it is not tested here.
  
      AdaScale doesn't have distributed state. Otherwise, it will need
      a unit test for checkpointing with DDP.
      """
      # Constants.
      num_grads_to_accum = 3
      in_dim = 5
  
      # Setup.
      def make_model_and_optim():
          model = Linear(in_dim, 2, bias=False)
          model = model.cuda()
          optim = AdaScale(
              SGD(model.parameters(), lr=0.1, momentum=0.9),
              num_gradients_to_accumulate=num_grads_to_accum,
          )
          return model, optim
  
      model, optim = make_model_and_optim()
  
      # Run a bit.
      def run_a_bit(replay_data=None):
          data = []
          replay_data_idx = 0
          for _ in range(6):  # run some steps
              for i in range(num_grads_to_accum):
                  if replay_data is None:
                      in_data = torch.rand(in_dim).cuda()
                      data.append(in_data)
                  else:
                      in_data = replay_data[replay_data_idx]
                      replay_data_idx += 1
                  out = model(in_data)
                  out.sum().backward()
                  if i == num_grads_to_accum - 1:
                      optim.step()
                      optim.zero_grad()
          return out, data
  
      run_a_bit()
  
      with tempfile.NamedTemporaryFile() as f:
          temp_file_name = f.name
  
          # Save a checkpoint.
          torch.save(
              {
                  "model": model.state_dict(),
                  "optim": optim.state_dict(),
              },
              temp_file_name,
          )
  
          # Train more.
          out, replay_data = run_a_bit()
  
          # Save the gain and out.
          expected_out = out.sum().item()
          expected_gain = optim.gain()
  
          # Load back the checkpoint.
          (
              model,
              optim,
          ) = make_model_and_optim()  # They both need to start afresh.
          ckpt = torch.load(temp_file_name)
          model.load_state_dict(ckpt["model"])
          optim.load_state_dict(ckpt["optim"])
  
          # Train the same steps.
          out, _ = run_a_bit(replay_data)
  
      # Assert the results.
      assert np.allclose(
          out.sum().item(), expected_out
      ), out.sum().item()
      assert np.allclose(optim.gain(), expected_gain), optim.gain()
  
imports: import tempfile
  import numpy
  import pytest
  import torch
  from torch import Tensor
  from torch.nn import Linear, Sequential
  from torch.optim import SGD
  from torch.optim.lr_scheduler import LambdaLR
  from fairscale.fair_dev.testing.golden_testing_data import adascale_test_data
  from fairscale.fair_dev.testing.testing import make_cudnn_deterministic, skip_if_no_cuda
  from fairscale.fair_dev.testing.testing_memory import find_tensor_by_shape
  from fairscale.optim import AdaScale
label_desc: Use image node_id 4 for calling a global function with example usage: test_state_checkpointing() and returns: model,  optim, out,  data