class_docstr: A mixin to share initialization of shared attributes and introduce prefetching.
class_docstr_tok: ["A","mixin","to","share","initialization","of","shared","attributes","and","introduce","prefetching","."]
code: def __init__(
      self,
      *,
      module: torch.nn.Module,
      device: Optional[torch.device] = None,
      precision: Optional[Union[str, torch.dtype]] = None,
      detect_anomaly: Optional[bool] = None,
      torch_compile_params: Optional[TorchCompileParams] = None,
  ) -> None:
      super().__init__()
      if torch_compile_params:
          _validate_torch_compile_available()
  
      self.device: torch.device = device or init_from_env()
      self.precision: Optional[torch.dtype] = (
          convert_precision_str_to_dtype(precision)
          if isinstance(precision, str)
          else precision
      )
  
      self.detect_anomaly = detect_anomaly
  
      # create autocast context based on precision and device type
      self.maybe_autocast_precision = torch.autocast(
          device_type=self.device.type,
          dtype=self.precision,
          enabled=self.precision is not None,
      )
  
      # cuda stream to use for moving data to device
      self._prefetch_stream: Optional[torch.cuda.streams.Stream] = (
          torch.cuda.Stream() if self.device.type == "cuda" else None
      )
      # dict mapping phase to whether the next batch which has been prefetched for that phase and is ready to be used
      self._phase_to_next_batch: dict[ActivePhase, Optional[TData]] = {
          ActivePhase.TRAIN: None,
          ActivePhase.EVALUATE: None,
          ActivePhase.PREDICT: None,
      }
  
      # dict mapping phase to whether the next batch for that phase has been prefetched and is ready to be used
      self._phase_to_prefetched: dict[ActivePhase, bool] = {
          ActivePhase.TRAIN: False,
          ActivePhase.EVALUATE: False,
          ActivePhase.PREDICT: False,
      }
      # whether the current batch is the last train batch
      self._is_last_batch: bool = False
  
imports: import contextlib
  from abc import ABCMeta, abstractmethod
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Any, Callable, Generic, Iterator, List, Optional, Tuple, TypeVar, Union
  import torch
  from pyre_extensions import none_throws
  from torch.cuda.amp import GradScaler
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim.swa_utils import SWALR
  from torchtnt.framework._loop_utils import _step_requires_iterator
  from torchtnt.framework.state import ActivePhase, EntryPoint, State
  from torchtnt.framework.unit import EvalUnit, PredictUnit, TPredictData, TrainUnit
  from torchtnt.framework.utils import get_timing_context
  from torchtnt.utils.device import copy_data_to_device, record_data_in_stream
  from torchtnt.utils.env import init_from_env
  from torchtnt.utils.lr_scheduler import TLRScheduler
  from torchtnt.utils.precision import convert_precision_str_to_dtype, get_grad_scaler_from_precision
  from torchtnt.utils.prepare_module import _is_fsdp_module, ActivationCheckpointParams, FSDPStrategy, prepare_fsdp, prepare_module, Strategy, TorchCompileParams
  from torchtnt.utils.swa import AveragedModel
  from torchtnt.utils.version import is_torch_version_ge_1_13_1
  from typing_extensions import Literal
label_desc: Use image node_id 1 to create a new _AutoUnitMixin object from inherited base classes: Generic with example: obj = _AutoUnitMixin()