class_docstr: Like :class:`SimpleTrainer`, but uses PyTorch's native automatic mixed precision
  in the training loop.
class_docstr_tok: ["Like",":","class",":","`","SimpleTrainer","`",",","but","uses","PyTorch","'s","native","automatic","mixed","precision","in","the","training","loop","."]
docstr: Implement the AMP training logic.
docstr_tok: ["Implement","the","AMP","training","logic","."]
code: def run_step(self):
      """
      Implement the AMP training logic.
      """
      assert (
          self.model.training
      ), "[AMPTrainer] model was changed to eval mode!"
      assert (
          torch.cuda.is_available()
      ), "[AMPTrainer] CUDA is required for AMP training!"
      from torch.cuda.amp import autocast
  
      start = time.perf_counter()
      data = next(self._data_loader_iter)
      data_time = time.perf_counter() - start
  
      if self.zero_grad_before_forward:
          self.optimizer.zero_grad()
      with autocast(dtype=self.precision):
          loss_dict = self.model(data)
          if isinstance(loss_dict, torch.Tensor):
              losses = loss_dict
              loss_dict = {"total_loss": loss_dict}
          else:
              losses = sum(loss_dict.values())
  
      if not self.zero_grad_before_forward:
          self.optimizer.zero_grad()
  
      self.grad_scaler.scale(losses).backward()
  
      if self.log_grad_scaler:
          storage = get_event_storage()
          storage.put_scalar(
              "[metric]grad_scaler", self.grad_scaler.get_scale()
          )
  
      self.after_backward()
  
      if self.async_write_metrics:
          # write metrics asynchronically
          self.concurrent_executor.submit(
              self._write_metrics, loss_dict, data_time, iter=self.iter
          )
      else:
          self._write_metrics(loss_dict, data_time)
  
      self.grad_scaler.step(self.optimizer)
      self.grad_scaler.update()
  
imports: import concurrent.futures
  import logging
  import numpy
  import time
  import weakref
  from typing import List, Mapping, Optional
  import torch
  from torch.nn.parallel import DataParallel, DistributedDataParallel
  import detectron2.utils.comm
  from detectron2.utils.events import EventStorage, get_event_storage
  from detectron2.utils.logger import _log_api_usage
label_desc: Use image node_id 2 for calling the AMPTrainer obj's underlying member method code with example usage: obj.run_step() without return types