code: def train_eval_train_test(self, module, t):
      model = module(t).cuda()
      optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
  
      _amp_state.allow_incoming_model_not_fp32 = True
      model, optimizer = amp.initialize(
          model, optimizer, opt_level="O1", verbosity=0
      )
      _amp_state.allow_incoming_model_not_fp32 = False
  
      def training_step():
          for param in model.parameters():
              param.grad = None
  
          loss = model(self.x).sum()
          _amp_state.loss_scalers[0]._loss_scale = 4.0
          with amp.scale_loss(loss, optimizer) as scaled_loss:
              scaled_loss.backward()
  
          self.assertEqual(
              len(
                  [
                      p.grad
                      for p in model.parameters()
                      if p.grad is not None
                  ]
              ),
              1,
          )
          self.assertEqual(
              model.weight.grad.type(), model.weight.type()
          )
  
          reference_grad = get_reference_grad(
              self.x, model.weight, model.ops
          )
  
          # Currently there's no difference in the allclose calls, so no need for branching,
          # but I'm keeping this in case we want different tolerances for fp16 and fp32 checks.
          if model.weight.grad.type() == "torch.cuda.HalfTensor":
              torch.testing.assert_close(
                  model.weight.grad.float(), reference_grad
              )
          elif model.weight.grad.type() == "torch.cuda.FloatTensor":
              torch.testing.assert_close(
                  model.weight.grad.float(), reference_grad
              )
          else:
              raise RuntimeError(
                  "model.weight.grad.type = {}".format(
                      model.weight.grad.type()
                  )
              )
  
          model.weight.data -= 1.0
  
      # Simulates first epoch
      training_step()
  
      # Simulates eval
      with torch.no_grad():
          loss = model(self.x).sum()
  
      # Simulates resuming training after eval
      training_step()
  
      _amp_state.handle._deactivate()
  
imports: import unittest
  import functools
  import itertools
  from apex import amp
  from apex.amp import _amp_state
  import torch
  from torch import nn
  import torch.nn.functional
  from utils import common_init, HALF, FLOAT, ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
label_desc: Use image node_id 3 for calling the TestCache obj's underlying member method code with example usage: obj.train_eval_train_test(module, t) without return types