class_docstr: Contains runtime information for any Determined workflow that uses the ``PyTorch`` API.
  
  With this class, users can do the following things:
  
  1. Wrap PyTorch models, optimizers, and LR schedulers with their Determined-compatible
     counterparts using :meth:`wrap_model`, :meth:`wrap_optimizer`, :meth:`wrap_lr_scheduler`,
     respectively. The Determined-compatible objects are capable of transparent
     distributed training, checkpointing and exporting, mixed-precision training,
     and gradient aggregation.
  2. Configure apex amp by calling :meth:`configure_apex_amp` (optional).
  3. Calculate the gradients with :meth:`backward` on a specified loss.
  4. Run an optimization step with :meth:`step_optimizer`.
  5. Functionalities inherited from :class:`determined.TrialContext`, including getting
     the runtime information and properly handling training data in distributed training.
class_docstr_tok: ["Contains","runtime","information","for","any","Determined","workflow","that","uses","the","``","PyTorch","``","API",".","With","this","class",",","users","can","do","the","following","things",":","1",".","Wrap","PyTorch","models",",","optimizers",",","and","LR","schedulers","with","their","Determined-compatible","counterparts","using",":","meth",":","`","wrap_model","`",",",":","meth",":","`","wrap_optimizer","`",",",":","meth",":","`","wrap_lr_scheduler","`",",","respectively",".","The","Determined-compatible","objects","are","capable","of","transparent","distributed","training",",","checkpointing","and","exporting",",","mixed-precision","training",",","and","gradient","aggregation",".","2",".","Configure","apex","amp","by","calling",":","meth",":","`","configure_apex_amp","`","(","optional",")",".","3",".","Calculate","the","gradients","with",":","meth",":","`","backward","`","on","a","specified","loss",".","4",".","Run","an","optimization","step","with",":","meth",":","`","step_optimizer","`",".","5",".","Functionalities","inherited","from",":","class",":","`","determined.TrialContext","`",",","including","getting","the","runtime","information","and","properly","handling","training","data","in","distributed","training","."]
docstr: Configure automatic mixed precision for your models and optimizers using NVIDIA's Apex
  PyTorch extension. Note that details for ``apex.amp`` are handled automatically within
  Determined after this call.
  
  This function must be called **after** you have finished constructing your models and
  optimizers with :meth:`wrap_model` and :meth:`wrap_optimizer`.
  
  This function has the same arguments as
  `apex.amp.initialize <https://nvidia.github.io/apex/amp.html#apex.amp.initialize>`_.
  
  .. warning::
      When using distributed training and automatic mixed precision,
      we only support ``num_losses=1`` and calling backward on the loss once.
  
  Arguments:
      models (``torch.nn.Module`` or list of ``torch.nn.Module`` s):  Model(s) to modify/cast.
      optimizers (``torch.optim.Optimizer`` or list of ``torch.optim.Optimizer`` s):
          Optimizers to modify/cast. REQUIRED for training.
      enabled (bool, optional, default=True):  If False, renders all Amp calls no-ops,
          so your script should run as if Amp were not present.
      opt_level (str, optional, default="O1"):  Pure or mixed precision optimization level.
          Accepted values are "O0", "O1", "O2", and "O3", explained in detail above.
      cast_model_type (``torch.dtype``, optional, default=None):  Optional property override,
          see above.
      patch_torch_functions (bool, optional, default=None):  Optional property override.
      keep_batchnorm_fp32 (bool or str, optional, default=None):  Optional property override.
          If passed as a string, must be the string "True" or "False".
      master_weights (bool, optional, default=None):  Optional property override.
      loss_scale (float or str, optional, default=None):  Optional property override.
          If passed as a string, must be a string representing a number, e.g., "128.0",
          or the string "dynamic".
      cast_model_outputs (torch.dtype, optional, default=None):  Option to ensure that
          the outputs of your model is always cast to a particular type regardless of
          ``opt_level``.
      num_losses (int, optional, default=1):  Option to tell Amp in advance how many
          losses/backward passes you plan to use.  When used in conjunction with the
          ``loss_id`` argument to ``amp.scale_loss``, enables Amp to use a different
          loss scale per loss/backward pass, which can improve stability.
          If ``num_losses`` is left to 1, Amp will still support multiple losses/backward
          passes, but use a single global loss scale for all of them.
      verbosity (int, default=1):  Set to 0 to suppress Amp-related output.
      min_loss_scale (float, default=None):  Sets a floor for the loss scale values that
          can be chosen by dynamic loss scaling.  The default value of None means that no
          floor is imposed. If dynamic loss scaling is not used, ``min_loss_scale`` is
          ignored.
      max_loss_scale (float, default=2.**24):  Sets a ceiling for the loss scale values
          that can be chosen by dynamic loss scaling.  If dynamic loss scaling is not used,
          ``max_loss_scale`` is ignored.
  
  Returns:
      Model(s) and optimizer(s) modified according to the ``opt_level``.
      If  ``optimizers`` args were lists, the corresponding return value will
      also be a list.
code: def configure_apex_amp(
      self,
      models: Union[torch.nn.Module, List[torch.nn.Module]],
      optimizers: Union[
          torch.optim.Optimizer, List[torch.optim.Optimizer]
      ],
      enabled: Optional[bool] = True,
      opt_level: Optional[str] = "O1",
      cast_model_type: Optional[torch.dtype] = None,
      patch_torch_functions: Optional[bool] = None,
      keep_batchnorm_fp32: Optional[Union[bool, str]] = None,
      master_weights: Optional[bool] = None,
      loss_scale: Optional[Union[float, str]] = None,
      cast_model_outputs: Optional[torch.dtype] = None,
      num_losses: Optional[int] = 1,
      verbosity: Optional[int] = 1,
      min_loss_scale: Optional[float] = None,
      max_loss_scale: Optional[float] = 2.0**24,
  ) -> Tuple:
      """
      Configure automatic mixed precision for your models and optimizers using NVIDIA's Apex
      PyTorch extension. Note that details for ``apex.amp`` are handled automatically within
      Determined after this call.
  
      This function must be called **after** you have finished constructing your models and
      optimizers with :meth:`wrap_model` and :meth:`wrap_optimizer`.
  
      This function has the same arguments as
      `apex.amp.initialize <https://nvidia.github.io/apex/amp.html#apex.amp.initialize>`_.
  
      .. warning::
          When using distributed training and automatic mixed precision,
          we only support ``num_losses=1`` and calling backward on the loss once.
  
      Arguments:
          models (``torch.nn.Module`` or list of ``torch.nn.Module`` s):  Model(s) to modify/cast.
          optimizers (``torch.optim.Optimizer`` or list of ``torch.optim.Optimizer`` s):
              Optimizers to modify/cast. REQUIRED for training.
          enabled (bool, optional, default=True):  If False, renders all Amp calls no-ops,
              so your script should run as if Amp were not present.
          opt_level (str, optional, default="O1"):  Pure or mixed precision optimization level.
              Accepted values are "O0", "O1", "O2", and "O3", explained in detail above.
          cast_model_type (``torch.dtype``, optional, default=None):  Optional property override,
              see above.
          patch_torch_functions (bool, optional, default=None):  Optional property override.
          keep_batchnorm_fp32 (bool or str, optional, default=None):  Optional property override.
              If passed as a string, must be the string "True" or "False".
          master_weights (bool, optional, default=None):  Optional property override.
          loss_scale (float or str, optional, default=None):  Optional property override.
              If passed as a string, must be a string representing a number, e.g., "128.0",
              or the string "dynamic".
          cast_model_outputs (torch.dtype, optional, default=None):  Option to ensure that
              the outputs of your model is always cast to a particular type regardless of
              ``opt_level``.
          num_losses (int, optional, default=1):  Option to tell Amp in advance how many
              losses/backward passes you plan to use.  When used in conjunction with the
              ``loss_id`` argument to ``amp.scale_loss``, enables Amp to use a different
              loss scale per loss/backward pass, which can improve stability.
              If ``num_losses`` is left to 1, Amp will still support multiple losses/backward
              passes, but use a single global loss scale for all of them.
          verbosity (int, default=1):  Set to 0 to suppress Amp-related output.
          min_loss_scale (float, default=None):  Sets a floor for the loss scale values that
              can be chosen by dynamic loss scaling.  The default value of None means that no
              floor is imposed. If dynamic loss scaling is not used, ``min_loss_scale`` is
              ignored.
          max_loss_scale (float, default=2.**24):  Sets a ceiling for the loss scale values
              that can be chosen by dynamic loss scaling.  If dynamic loss scaling is not used,
              ``max_loss_scale`` is ignored.
  
      Returns:
          Model(s) and optimizer(s) modified according to the ``opt_level``.
          If  ``optimizers`` args were lists, the corresponding return value will
          also be a list.
      """
      if not enabled or not self._managed_training:
          return models, optimizers
  
      if self._scaler is not None and self._scaler.is_enabled():
          raise det.errors.InvalidExperimentException(
              "Do not mix APEX with PyTorch AMP."
          )
  
      warnings.warn(
          "PyTorchTrial support for NVIDIA/apex has been deprecated and will be removed "
          "in a future version. We recommend users to migrate to Torch AMP (`torch.cuda.amp`).",
          FutureWarning,
          stacklevel=2,
      )
  
      if self._use_apex:
          raise det.errors.InvalidExperimentException(
              "Please only call configure_apex_amp once."
          )
  
      if self.distributed.size > 1:
          if num_losses != 1:
              raise det.errors.InvalidExperimentException(
                  "When using distributed training, "
                  "Determined only supports configure_apex_amp with num_losses = 1.",
              )
          if self._aggregation_frequency > 1:
              raise det.errors.InvalidExperimentException(
                  "context.configure_apex_amp is not supported with "
                  "distributed training and "
                  "aggregation frequency > 1.",
              )
  
      if not torch.cuda.is_available():
          raise det.errors.InvalidExperimentException(
              "context.configure_apex_amp is supported only on GPU slots.",
          )
  
      self._use_apex = True
  
      if self._distributed_backend.use_torch():
          # We need to get the pre-wrapped input models to initialize APEX because
          if isinstance(models, list):
              models = [
                  self._wrapped_models[wrapped_model]
                  for wrapped_model in models
              ]
          else:
              models = self._wrapped_models[models]
  
      logger.info(
          f"Enabling mixed precision training with opt_level: {opt_level}."
      )
      models, optimizers = apex.amp.initialize(
          models=models,
          optimizers=optimizers,
          enabled=enabled,
          opt_level=opt_level,
          cast_model_type=cast_model_type,
          patch_torch_functions=patch_torch_functions,
          keep_batchnorm_fp32=keep_batchnorm_fp32,
          master_weights=master_weights,
          loss_scale=loss_scale,
          cast_model_outputs=cast_model_outputs,
          num_losses=num_losses,
          min_loss_scale=min_loss_scale,
          max_loss_scale=max_loss_scale,
          verbosity=verbosity
          if self.distributed.get_rank() == 0 or self._debug_enabled
          else 0,
      )
  
      if not isinstance(models, list):
          self.models = [models]
  
      if (
          self.distributed.size > 1
          and self._distributed_backend.use_torch()
      ):
          # If Torch DDP is in use, re-wrap the models
          self.models = [
              self._PyTorchDistributedDataParallel(model)
              for model in self.models
          ]
  
      if not isinstance(optimizers, list):
          self.optimizers = [optimizers]
      return models, optimizers
  
imports: import contextlib
  import logging
  import pathlib
  import warnings
  from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union
  import torch
  from torch import nn
  import determined
  from determined import profiler, pytorch, util
  from determined.horovod import hvd
label_desc: Use image node_id 25 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.configure_apex_amp(models, optimizers, enabled, opt_level, cast_model_type, patch_torch_functions, keep_batchnorm_fp32, master_weights, loss_scale, cast_model_outputs, num_losses, verbosity, min_loss_scale, max_loss_scale) and returns: models,  optimizers, models,  optimizers