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: Prepares to use automatic mixed precision through PyTorch’s native AMP API. The returned
  scaler should be passed to ``step_optimizer``, but usage does not otherwise differ from
  vanilla PyTorch APIs. Loss should be scaled before calling ``backward``, ``unscale_`` should
  be called before clipping gradients, ``update`` should be called after stepping all
  optimizers, etc.
  
  PyTorch 1.6 or greater is required for this feature.
  
  Arguments:
      scaler (``torch.cuda.amp.GradScaler``):  Scaler to wrap and track.
  
  Returns:
      The scaler. It may be wrapped to add additional functionality for use in Determined.
code: def wrap_scaler(self, scaler: Any) -> Any:
      """
      Prepares to use automatic mixed precision through PyTorch’s native AMP API. The returned
      scaler should be passed to ``step_optimizer``, but usage does not otherwise differ from
      vanilla PyTorch APIs. Loss should be scaled before calling ``backward``, ``unscale_`` should
      be called before clipping gradients, ``update`` should be called after stepping all
      optimizers, etc.
  
      PyTorch 1.6 or greater is required for this feature.
  
      Arguments:
          scaler (``torch.cuda.amp.GradScaler``):  Scaler to wrap and track.
  
      Returns:
          The scaler. It may be wrapped to add additional functionality for use in Determined.
      """
  
      if not HAVE_AMP:
          raise det.errors.InvalidExperimentException(
              "Using context.wrap_scaler() requires PyTorch >= 1.6.",
          )
  
      if self._use_apex:
          raise det.errors.InvalidExperimentException(
              "Do not mix APEX with PyTorch AMP."
          )
  
      if self._scaler is not None:
          raise det.errors.InvalidExperimentException(
              "Please only call wrap_scaler or use_amp once.",
          )
  
      if self.models:
          raise det.errors.InvalidExperimentException(
              "Please call wrap_scaler before wrap_model.",
          )
  
      # We don't need to check if CUDA is available because if it is not, a GradScaler is
      #  disabled when initialized, and we allow for disabled scalers to exist.
  
      self._scaler = scaler
  
      return scaler
  
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 24 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.wrap_scaler(scaler) and returns: scaler