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 autocast_forward_pass(
      self, to_wrap: torch.nn.Module
  ) -> torch.nn.Module:
      # First, ensure the forward pass is wrapped in an autocast context:
      class _AutocastForwardPassModel(type(to_wrap)):  # type: ignore
          def __init__(wrapper) -> None:
              self.model = to_wrap
  
          def __getattr__(wrapper, name):  # type: ignore
              return getattr(to_wrap, name)
  
          def __setattr__(wrapper, name, value):  # type: ignore
              return setattr(to_wrap, name, value)
  
          def __delattr__(wrapper, name):  # type: ignore
              return delattr(to_wrap, name)
  
          def forward(wrapper, *arg, **kwarg):  # type: ignore
              with amp.autocast():
                  return to_wrap.forward(*arg, **kwarg)
  
      wrapped = _AutocastForwardPassModel()
  
      # To not print errors recursively or every forward pass
      warned_types = set()
  
      # Second, eliminate any need for the loss functions to be in that context:
      def end_fp16(
          module: torch.nn.Module, input: Any, output: Any
      ) -> Any:  # noqa: A002
          if isinstance(output, torch.Tensor):
              return (
                  output.float()
                  if output.dtype == torch.float16
                  else output
              )
          if isinstance(output, dict):
              return {
                  k: end_fp16(module, input, v)
                  for k, v in output.items()
              }
          if isinstance(output, list):
              return [end_fp16(module, input, d) for d in output]
          if isinstance(output, tuple):
              return tuple(end_fp16(module, input, d) for d in output)
          # If there are other types that embed Tensors still using fp16 and loss computation
          # subsequently fails, then experimental.use_amp should not be used and the forward pass
          # should be manually wrapped in an autocast context.
          if type(output) not in warned_types:
              warned_types.add(type(output))
              logger.warn(
                  f"Unexpected type '{type(output).__name__}' outputted by model in experimental "
                  "AMP mode."
              )
          return output
  
      wrapped.register_forward_hook(end_fp16)
  
      return wrapped
  
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 11 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.autocast_forward_pass(to_wrap) and returns: wrapped, output, getattr, setattr, delattr, unknown, unknown, unknown, tuple, to_wrap