docstr: AutoPredictUnit is a convenience for users who are running inference and would like to have certain features handled for them, such as:
  - Moving data to the correct device.
  - Running inference under a mixed precision context.
  - Handling data parallel replication, especially if the module cannot fit on a single device using FullyShardedDataParallel.
  - Profiling the data transfer to device and forward pass.
  - Interleaving moving the next batch to the device with running the module's forward pass on the current batch.
  
  Additionally, the AutoPredictUnit offers an optional hook ``on_predict_step_end`` to further post-process module outputs if desired.
  
  Then use with the :py:func:`~torchtnt.framework.predict` entry point.
  
  For more advanced customization, directly use the :class:`~torchtnt.framework.unit.PredictUnit` interface.
  
  Args:
      module: module to be used during prediction.
      device: the device to be used.
      precision: the precision to use in training, as either a string or a torch.dtype.
      strategy: the data parallelization strategy to be used. if a string, must be one of ``ddp`` or ``fsdp``.
      torch_compile_params: params for Torch compile https://pytorch.org/docs/stable/generated/torch.compile.html
      detect_anomaly: whether to enable anomaly detection for the autograd engine https://pytorch.org/docs/stable/autograd.html#anomaly-detection
  
  Note:
      Torch compile support is only available in PyTorch 2.0 or higher.
docstr_tok: ["AutoPredictUnit","is","a","convenience","for","users","who","are","running","inference","and","would","like","to","have","certain","features","handled","for","them",",","such","as",":","-","Moving","data","to","the","correct","device",".","-","Running","inference","under","a","mixed","precision","context",".","-","Handling","data","parallel","replication",",","especially","if","the","module","can","not","fit","on","a","single","device","using","FullyShardedDataParallel",".","-","Profiling","the","data","transfer","to","device","and","forward","pass",".","-","Interleaving","moving","the","next","batch","to","the","device","with","running","the","module","'s","forward","pass","on","the","current","batch",".","Additionally",",","the","AutoPredictUnit","offers","an","optional","hook","``","on_predict_step_end","``","to","further","post-process","module","outputs","if","desired",".","Then","use","with","the",":","py",":","func",":","`","~torchtnt.framework.predict","`","entry","point",".","For","more","advanced","customization",",","directly","use","the",":","class",":","`","~torchtnt.framework.unit.PredictUnit","`","interface",".","Args",":","module",":","module","to","be","used","during","prediction",".","device",":","the","device","to","be","used",".","precision",":","the","precision","to","use","in","training",",","as","either","a","string","or","a","torch.dtype",".","strategy",":","the","data","parallelization","strategy","to","be","used",".","if","a","string",",","must","be","one","of","``","ddp","``","or","``","fsdp","``",".","torch_compile_params",":","params","for","Torch","compile","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.compile.html","detect_anomaly",":","whether","to","enable","anomaly","detection","for","the","autograd","engine","https",":","\/\/pytorch.org\/docs\/stable\/autograd.html","#","anomaly-detection","Note",":","Torch","compile","support","is","only","available","in","PyTorch","2.0","or","higher","."]
code: def __init__(
      self,
      *,
      module: torch.nn.Module,
      device: Optional[torch.device] = None,
      strategy: Optional[Union[Strategy, str]] = None,
      precision: Optional[Union[str, torch.dtype]] = None,
      torch_compile_params: Optional[TorchCompileParams] = None,
      detect_anomaly: Optional[bool] = None,
  ) -> None:
      """
      AutoPredictUnit is a convenience for users who are running inference and would like to have certain features handled for them, such as:
      - Moving data to the correct device.
      - Running inference under a mixed precision context.
      - Handling data parallel replication, especially if the module cannot fit on a single device using FullyShardedDataParallel.
      - Profiling the data transfer to device and forward pass.
      - Interleaving moving the next batch to the device with running the module's forward pass on the current batch.
  
      Additionally, the AutoPredictUnit offers an optional hook ``on_predict_step_end`` to further post-process module outputs if desired.
  
      Then use with the :py:func:`~torchtnt.framework.predict` entry point.
  
      For more advanced customization, directly use the :class:`~torchtnt.framework.unit.PredictUnit` interface.
  
      Args:
          module: module to be used during prediction.
          device: the device to be used.
          precision: the precision to use in training, as either a string or a torch.dtype.
          strategy: the data parallelization strategy to be used. if a string, must be one of ``ddp`` or ``fsdp``.
          torch_compile_params: params for Torch compile https://pytorch.org/docs/stable/generated/torch.compile.html
          detect_anomaly: whether to enable anomaly detection for the autograd engine https://pytorch.org/docs/stable/autograd.html#anomaly-detection
  
      Note:
          Torch compile support is only available in PyTorch 2.0 or higher.
      """
      super().__init__(
          module=module,
          device=device,
          precision=precision,
          torch_compile_params=torch_compile_params,
          detect_anomaly=detect_anomaly,
      )
      self.module: torch.nn.Module = prepare_module(
          module,
          self.device,
          strategy=strategy,
          torch_compile_params=torch_compile_params,
      )
  
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 AutoPredictUnit object from inherited base classes: _AutoUnitMixin, PredictUnit with example: obj = AutoPredictUnit()