code: def test_restoring(self):
      nb_epochs = 10
      nb_epochs_restore = nb_epochs // 2
      for opt_level in self.test_opt_levels:
          for res_opt_level in self.test_opt_levels:
              for amp_before_load in [True, False]:
                  for num_losses in range(1, 3):
                      test_setup = (
                          "#" * 75
                          + "\n"
                          + f"opt_level {opt_level}\n"
                          + f"restore_opt_level {res_opt_level}\n"
                          + f"amp_before_load {amp_before_load}\n"
                          + f"num_losses {num_losses}\n"
                      )
  
                      self.seed()
  
                      # Create reference model
                      model = MyModel().to("cuda")
  
                      optimizer = optim.SGD(
                          model.parameters(), lr=self.initial_lr
                      )
  
                      # Initialize with num_losses*2 for the original model and the restored one
                      model, optimizer = amp.initialize(
                          model,
                          optimizer,
                          opt_level=opt_level,
                          num_losses=num_losses * 2,
                          verbosity=0,
                      )
  
                      # Compare training behavior for same restore option
                      # We cannot really generalize it, since a saved model in O0
                      # would introduce a skipped step in O1, which will raise an error
                      if opt_level == res_opt_level:
                          # train for nb_epochs and restore after nb_epochs_restore
                          for epoch in range(nb_epochs):
                              x = torch.randn(
                                  16, 3, 24, 24, device="cuda"
                              )
                              output = self.train_step(
                                  model, optimizer, x, range(num_losses)
                              )
                              # Initialize model one step before comparing.
                              # Otherwise the batchnorm layers will be updated
                              # additionally in restore_model
                              if epoch == (nb_epochs_restore - 1):
                                  # Load model and optimizer
                                  checkpoint = {
                                      "model": model.state_dict(),
                                      "optimizer": optimizer.state_dict(),
                                      "amp": amp.state_dict(),
                                  }
                                  # Check state_dict for FP32 tensors
                                  self.check_state_dict_fp32(
                                      checkpoint["model"]
                                  )
  
                                  # Restore model
                                  restore_model = MyModel().to("cuda")
                                  restore_optimizer = optim.SGD(
                                      restore_model.parameters(),
                                      lr=self.initial_lr,
                                  )
  
                                  if amp_before_load:
                                      (
                                          restore_model,
                                          restore_optimizer,
                                      ) = amp.initialize(
                                          restore_model,
                                          restore_optimizer,
                                          opt_level=res_opt_level,
                                          num_losses=num_losses * 2,
                                          verbosity=0,
                                      )
  
                                  restore_model.load_state_dict(
                                      checkpoint["model"]
                                  )
                                  restore_optimizer.load_state_dict(
                                      checkpoint["optimizer"]
                                  )
                                  # FIXME: We cannot test the amp.state_dict in the same script
                                  # amp.load_state_dict(checkpoint['amp'])
  
                                  if not amp_before_load:
                                      (
                                          restore_model,
                                          restore_optimizer,
                                      ) = amp.initialize(
                                          restore_model,
                                          restore_optimizer,
                                          opt_level=res_opt_level,
                                          num_losses=num_losses * 2,
                                          verbosity=0,
                                      )
  
                              elif epoch >= nb_epochs_restore:
                                  restore_output = self.train_step(
                                      restore_model,
                                      restore_optimizer,
                                      x,
                                      range(num_losses, num_losses * 2),
                                  )
                                  self.assertTrue(
                                      torch.allclose(
                                          output.float(),
                                          restore_output.float(),
                                      ),
                                      "Output of reference and restored models differ for "
                                      + test_setup,
                                  )
                                  self.compare_models(
                                      model, restore_model, test_setup
                                  )
                      # if opt_level != res_opt_level
                      else:
                          # skip tests for different opt_levels
                          continue
  
imports: import unittest
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from apex import amp
  from utils import common_init, FLOAT
label_desc: Use image node_id 6 for calling the TestCheckpointing obj's underlying member method code with example usage: obj.test_restoring() without return types