class_docstr: Plugin for Automatic Mixed Precision (AMP) training with ``torch.autocast``.
  
  Args:
      precision: Whether to use ``torch.float16`` (``'16-mixed'``) or ``torch.bfloat16`` (``'bf16-mixed'``).
      device: The device for ``torch.autocast``.
      scaler: An optional :class:`torch.cuda.amp.GradScaler` to use.
code: def optimizer_step(  # type: ignore[override]
      self,
      optimizer: Optimizable,
      model: "pl.LightningModule",
      closure: Callable[[], Any],
      **kwargs: Any,
  ) -> Any:
      if self.scaler is None:
          # skip scaler logic, as bfloat16 does not require scaler
          return super().optimizer_step(
              optimizer, model=model, closure=closure, **kwargs
          )
      if isinstance(optimizer, LBFGS):
          raise MisconfigurationException(
              "AMP and the LBFGS optimizer are not compatible."
          )
      closure_result = closure()
  
      # If backward was skipped in automatic optimization (return None), unscaling is not needed
      skip_unscaling = (
          closure_result is None and model.automatic_optimization
      )
  
      if (
          not _optimizer_handles_unscaling(optimizer)
          and not skip_unscaling
      ):
          # Unscaling needs to be performed here in case we are going to apply gradient clipping.
          # Optimizers that perform unscaling in their `.step()` method are not supported (e.g., fused Adam).
          # Note: `unscale` happens after the closure is executed, but before the `on_before_optimizer_step` hook.
          self.scaler.unscale_(optimizer)
  
      self._after_closure(model, optimizer)
  
      # in manual optimization, the closure does not return a value
      if not skip_unscaling:
          # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
          step_output = self.scaler.step(optimizer, **kwargs)
          self.scaler.update()
          return step_output
      return closure_result
  
imports: from contextlib import contextmanager
  from typing import Any, Callable, Dict, Generator, Literal, Optional, Union
  import torch
  from torch import Tensor
  from torch.optim import LBFGS, Optimizer
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.accelerators.cuda import _patch_cuda_is_available
  from lightning.fabric.plugins.precision.amp import _optimizer_handles_unscaling
  from lightning.fabric.utilities.types import Optimizable
  from lightning.pytorch.plugins.precision.precision import Precision
  from lightning.pytorch.utilities import GradClipAlgorithmType
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
label_desc: Use image node_id 3 for calling the MixedPrecision obj's underlying member method code with example usage: obj.optimizer_step(optimizer, model, closure) and returns: closure_result, super, step_output