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","."]
code: def __init__(
      self,
      core_context: det.core.Context,
      trial_seed: Optional[int],
      hparams: Optional[Dict],
      slots_per_trial: int,
      num_gpus: int,
      exp_conf: Optional[Dict[str, Any]],
      aggregation_frequency: int,
      steps_completed: int,
      managed_training: bool,
      debug_enabled: bool,
      enable_tensorboard_logging: bool = True,
  ) -> None:
      self._core = core_context
      self.distributed = self._core.distributed
      pytorch._PyTorchReducerContext.__init__(
          self, self.distributed.allgather
      )
      self._per_slot_batch_size, self._global_batch_size = (
          util.calculate_batch_sizes(
              hparams=hparams,
              slots_per_trial=slots_per_trial,
              trialname="PyTorchTrial",
          )
          if hparams and hparams.get("global_batch_size", None)
          else (None, None)
      )
      self._hparams = hparams
      self._num_gpus = num_gpus
      self._debug_enabled = debug_enabled
      self._exp_conf = exp_conf
  
      self._trial_seed = trial_seed
      self._distributed_backend = det._DistributedBackend()
      self._steps_completed = steps_completed
      self.device = self._init_device()
  
      # Track which types we have issued warnings for in to_device().
      self._to_device_warned_types = set()  # type: Set[Type]
  
      # The following attributes are initialized during the lifetime of
      # a PyTorchTrialContext.
      self.models = []  # type: List[nn.Module]
      self.optimizers = []  # type: List[torch.optim.Optimizer]
      self.profiler = None  # type: Any
      self.lr_schedulers = []  # type: List[pytorch.LRScheduler]
      self._epoch_len = None  # type: Optional[int]
  
      # Keep a map of wrapped models to their original input forms, which is needed
      # by torch DDP and apex to initialize in the correct order
      self._wrapped_models = {}  # type: Dict[nn.Module, nn.Module]
  
      # Keep a map of optimizer configs set in wrap_optimizer which can differ per-optimizer
      self._optimizer_configs = {}  # type: Dict[Any, Any]
  
      # Use a main model to contain all the models because when using horovod
      # to broadcast the states of models we want to avoid name conflicts for these
      # states, so we set all the models to be submodule of the main model with
      # different names using __setattr__ and use the state_dict of the main model
      # for broadcasting. Note that broadcast_parameters only accepts state_dict()
      # although its doc says it also accepts named_parameters()
      self._main_model = nn.Module()
      self._scaler = None
      self._use_apex = False
      self._loss_ids = {}  # type: Dict[torch.Tensor, int]
      self._last_backward_batch_idx = None  # type: Optional[int]
      self._current_batch_idx = None  # type: Optional[int]
  
      self.experimental = pytorch.PyTorchExperimentalContext(self)
      self._reducers = pytorch._PyTorchReducerContext()
      self._determined_profiler = (
          None
      )  # type: Optional[profiler.ProfilerAgent]
  
      self._managed_training = managed_training
  
      self._aggregation_frequency = aggregation_frequency
  
      self._fp16_compression_default = False
      self._average_aggregated_gradients_default = True
      self._is_pre_trainer = False
  
      self._stop_requested = False
  
      self._tbd_writer = None  # type: Optional[Any]
      self._enable_tensorboard_logging = enable_tensorboard_logging
  
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 1 to create a new PyTorchTrialContext object from inherited base classes: pytorch with example: obj = PyTorchTrialContext(core_context, trial_seed, hparams, slots_per_trial, num_gpus, exp_conf, aggregation_frequency, steps_completed, managed_training, debug_enabled, enable_tensorboard_logging)