class_docstr: The AutoUnit is a convenience for users who are training with stochastic gradient descent and would like to have model optimization
  and data parallel replication handled for them.
  The AutoUnit subclasses :class:`~torchtnt.framework.unit.TrainUnit`, :class:`~torchtnt.framework.unit.EvalUnit`, and
  :class:`~torchtnt.framework.unit.PredictUnit` and implements the ``train_step``, ``eval_step``, and ``predict_step`` methods for the user.
  
  For the ``train_step`` it runs:
  
  - forward pass and loss computation
  - backward pass
  - optimizer step
  
  For the ``eval_step`` it only runs forward and loss computation.
  
  For the ``predict_step`` it only runs forward computation.
  
  To benefit from the AutoUnit, the user must subclass it and implement the ``compute_loss`` and ``configure_optimizers_and_lr_scheduler`` methods.
  Additionally, the AutoUnit offers these optional hooks:
  
  - ``on_train_step_end``
  - ``on_eval_step_end``
  - ``on_predict_step_end``
  
  Then use with the :py:func:`~torchtnt.framework.train`, :py:func:`~torchtnt.framework.evaluate`, :py:func:`~torchtnt.framework.fit`, or
  :py:func:`~torchtnt.framework.predict` entry point as normal.
  
  For more advanced customization, directly use the :class:`~torchtnt.framework.unit.TrainUnit`, :class:`~torchtnt.framework.unit.EvalUnit`,
  and :class:`~torchtnt.framework.unit.PredictUnit` interfaces.
  
  Args:
      module: module to be used during training/evaluation.
      device: the device to be used.
      strategy: the data parallelization strategy to be used. if a string, must be one of ``ddp`` or ``fsdp``.
      step_lr_interval: whether to step lr_scheduler every step or every epoch. Defaults to every epoch.
      precision: the precision to use in training/evaluation, as either a string or a torch.dtype.
      gradient_accumulation_steps: how many batches to accumulate gradients over.
      detect_anomaly: whether to enable anomaly detection for the autograd engine https://pytorch.org/docs/stable/autograd.html#anomaly-detection
      clip_grad_norm: max norm of the gradients for clipping https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
      clip_grad_value: max value of the gradients for clipping https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_value_.html
      swa_params: params for stochastic weight averaging https://pytorch.org/docs/stable/optim.html#stochastic-weight-averaging
      torch_compile_params: params for Torch compile https://pytorch.org/docs/stable/generated/torch.compile.html
      activation_checkpoint_params: params for enabling activation checkpointing
      training: if True, the optimizer and optionally LR scheduler will be created after the class is initialized.
  
  Note:
      If :class:`~torchtnt.utils.prepare_module.FSDPStrategy` and SWAParams are passed in, the swa model will be sharded with the same FSDP parameters.
  
  Note:
      Torch compile support is only available in PyTorch 2.0 or higher.
class_docstr_tok: ["The","AutoUnit","is","a","convenience","for","users","who","are","training","with","stochastic","gradient","descent","and","would","like","to","have","model","optimization","and","data","parallel","replication","handled","for","them",".","The","AutoUnit","subclasses",":","class",":","`","~torchtnt.framework.unit.TrainUnit","`",",",":","class",":","`","~torchtnt.framework.unit.EvalUnit","`",",","and",":","class",":","`","~torchtnt.framework.unit.PredictUnit","`","and","implements","the","``","train_step","``",",","``","eval_step","``",",","and","``","predict_step","``","methods","for","the","user",".","For","the","``","train_step","``","it","runs",":","-","forward","pass","and","loss","computation","-","backward","pass","-","optimizer","step","For","the","``","eval_step","``","it","only","runs","forward","and","loss","computation",".","For","the","``","predict_step","``","it","only","runs","forward","computation",".","To","benefit","from","the","AutoUnit",",","the","user","must","subclass","it","and","implement","the","``","compute_loss","``","and","``","configure_optimizers_and_lr_scheduler","``","methods",".","Additionally",",","the","AutoUnit","offers","these","optional","hooks",":","-","``","on_train_step_end","``","-","``","on_eval_step_end","``","-","``","on_predict_step_end","``","Then","use","with","the",":","py",":","func",":","`","~torchtnt.framework.train","`",",",":","py",":","func",":","`","~torchtnt.framework.evaluate","`",",",":","py",":","func",":","`","~torchtnt.framework.fit","`",",","or",":","py",":","func",":","`","~torchtnt.framework.predict","`","entry","point","as","normal",".","For","more","advanced","customization",",","directly","use","the",":","class",":","`","~torchtnt.framework.unit.TrainUnit","`",",",":","class",":","`","~torchtnt.framework.unit.EvalUnit","`",",","and",":","class",":","`","~torchtnt.framework.unit.PredictUnit","`","interfaces",".","Args",":","module",":","module","to","be","used","during","training\/evaluation",".","device",":","the","device","to","be","used",".","strategy",":","the","data","parallelization","strategy","to","be","used",".","if","a","string",",","must","be","one","of","``","ddp","``","or","``","fsdp","``",".","step_lr_interval",":","whether","to","step","lr_scheduler","every","step","or","every","epoch",".","Defaults","to","every","epoch",".","precision",":","the","precision","to","use","in","training\/evaluation",",","as","either","a","string","or","a","torch.dtype",".","gradient_accumulation_steps",":","how","many","batches","to","accumulate","gradients","over",".","detect_anomaly",":","whether","to","enable","anomaly","detection","for","the","autograd","engine","https",":","\/\/pytorch.org\/docs\/stable\/autograd.html","#","anomaly-detection","clip_grad_norm",":","max","norm","of","the","gradients","for","clipping","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.nn.utils.clip_grad_norm_.html","clip_grad_value",":","max","value","of","the","gradients","for","clipping","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.nn.utils.clip_grad_value_.html","swa_params",":","params","for","stochastic","weight","averaging","https",":","\/\/pytorch.org\/docs\/stable\/optim.html","#","stochastic-weight-averaging","torch_compile_params",":","params","for","Torch","compile","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.compile.html","activation_checkpoint_params",":","params","for","enabling","activation","checkpointing","training",":","if","True",",","the","optimizer","and","optionally","LR","scheduler","will","be","created","after","the","class","is","initialized",".","Note",":","If",":","class",":","`","~torchtnt.utils.prepare_module.FSDPStrategy","`","and","SWAParams","are","passed","in",",","the","swa","model","will","be","sharded","with","the","same","FSDP","parameters",".","Note",":","Torch","compile","support","is","only","available","in","PyTorch","2.0","or","higher","."]
docstr: The user should implement this method with their optimizer and learning rate scheduler construction code. This will be called upon initialization of
  the AutoUnit.
  
  Args:
      module: the module with which to construct optimizer and lr_scheduler
  
  Returns:
      A tuple containing optimizer and optionally the learning rate scheduler
docstr_tok: ["The","user","should","implement","this","method","with","their","optimizer","and","learning","rate","scheduler","construction","code",".","This","will","be","called","upon","initialization","of","the","AutoUnit",".","Args",":","module",":","the","module","with","which","to","construct","optimizer","and","lr_scheduler","Returns",":","A","tuple","containing","optimizer","and","optionally","the","learning","rate","scheduler"]
code: def configure_optimizers_and_lr_scheduler(
      self, module: torch.nn.Module
  ) -> Tuple[torch.optim.Optimizer, Optional[TLRScheduler]]:
      """
      The user should implement this method with their optimizer and learning rate scheduler construction code. This will be called upon initialization of
      the AutoUnit.
  
      Args:
          module: the module with which to construct optimizer and lr_scheduler
  
      Returns:
          A tuple containing optimizer and optionally the learning rate scheduler
      """
      ...
  
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 2 for calling the AutoUnit obj's underlying member method code with example usage: obj.configure_optimizers_and_lr_scheduler(module) without return types