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","."]
docstr: The user can override this method with custom code to copy data to device. This will be called at the start of every ``train_step``/``eval_step``/``predict_step``.
  By default this uses the utility function :py:func:`~torchtnt.utils.copy_data_to_device`.
  
  If on GPU, this method will be called on a separate CUDA stream.
  
  Args:
      state: a State object which is passed from the ``train_step``/``eval_step``/``predict_step``
      data: a batch of data which is passed from the ``train_step``/``eval_step``/``predict_step``
      non_blocking: parameter to pass to ``torch.tensor.to``
  
  Returns:
      A batch of data which is on the device
docstr_tok: ["The","user","can","override","this","method","with","custom","code","to","copy","data","to","device",".","This","will","be","called","at","the","start","of","every","``","train_step","``","\/","``","eval_step","``","\/","``","predict_step","``",".","By","default","this","uses","the","utility","function",":","py",":","func",":","`","~torchtnt.utils.copy_data_to_device","`",".","If","on","GPU",",","this","method","will","be","called","on","a","separate","CUDA","stream",".","Args",":","state",":","a","State","object","which","is","passed","from","the","``","train_step","``","\/","``","eval_step","``","\/","``","predict_step","``","data",":","a","batch","of","data","which","is","passed","from","the","``","train_step","``","\/","``","eval_step","``","\/","``","predict_step","``","non_blocking",":","parameter","to","pass","to","``","torch.tensor.to","``","Returns",":","A","batch","of","data","which","is","on","the","device"]
code: def move_data_to_device(
      self, state: State, data: TData, non_blocking: bool
  ) -> TData:
      """
      The user can override this method with custom code to copy data to device. This will be called at the start of every ``train_step``/``eval_step``/``predict_step``.
      By default this uses the utility function :py:func:`~torchtnt.utils.copy_data_to_device`.
  
      If on GPU, this method will be called on a separate CUDA stream.
  
      Args:
          state: a State object which is passed from the ``train_step``/``eval_step``/``predict_step``
          data: a batch of data which is passed from the ``train_step``/``eval_step``/``predict_step``
          non_blocking: parameter to pass to ``torch.tensor.to``
  
      Returns:
          A batch of data which is on the device
      """
      return copy_data_to_device(
          data, self.device, non_blocking=non_blocking
      )
  
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 _AutoUnitMixin obj's underlying member method code with example usage: obj.move_data_to_device(state, data, non_blocking) and returns: copy_data_to_device