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: Returns a wrapped LR scheduler.
  
  The LR scheduler must use an optimizer wrapped by :meth:`wrap_optimizer`.  If ``apex.amp``
  is in use, the optimizer must also have been configured with :meth:`configure_apex_amp`.
code: def wrap_lr_scheduler(
      self,
      lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
      step_mode: pytorch.LRScheduler.StepMode,
      frequency: int = 1,
  ) -> torch.optim.lr_scheduler._LRScheduler:
      """
      Returns a wrapped LR scheduler.
  
      The LR scheduler must use an optimizer wrapped by :meth:`wrap_optimizer`.  If ``apex.amp``
      is in use, the optimizer must also have been configured with :meth:`configure_apex_amp`.
      """
      if isinstance(
          lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau
      ):
          if step_mode != pytorch.LRScheduler.StepMode.MANUAL_STEP:
              raise det.errors.InvalidExperimentException(
                  "detected that context.wrap_lr_scheduler() was called with an instance of "
                  "torch.optim.lr_scheduler.ReduceLROnPlateau as the lr_scheduler.  This lr "
                  "scheduler class does not have the usual step() parameters, and so it can "
                  "only be used with step_mode=MANUAL_STEP.\n"
                  "\n"
                  "For example, if you wanted to step it on every validation step, you might "
                  "wrap your lr_scheduler and pass it to a callback like this:\n"
                  "\n"
                  "class MyLRStepper(PyTorchCallback):\n"
                  "    def __init__(self, wrapped_lr_scheduler):\n"
                  "        self.wrapped_lr_scheduler = wrapped_lr_scheduler\n"
                  "\n"
                  "    def on_validation_end(self, metrics):\n"
                  '        self.wrapped_lr_scheduler.step(metrics["validation_error"])\n'
              )
  
      opt = getattr(lr_scheduler, "optimizer", None)
      if opt is not None:
          if opt not in self.optimizers:
              raise det.errors.InvalidExperimentException(
                  "Must use an optimizer that is returned by wrap_optimizer().",
              )
      wrapped = pytorch.LRScheduler(lr_scheduler, step_mode, frequency)
      self.lr_schedulers.append(wrapped)
  
      # Return the original LR scheduler to the user in case they have customizations that we
      # don't care about.
      return lr_scheduler
  
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 14 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.wrap_lr_scheduler(lr_scheduler, step_mode, frequency) and returns: lr_scheduler