code: def train(
      train_unit: TTrainUnit,
      train_dataloader: Iterable[TTrainData],
      *,
      max_epochs: Optional[int] = None,
      max_steps: Optional[int] = None,
      max_steps_per_epoch: Optional[int] = None,
      callbacks: Optional[List[Callback]] = None,
      timer: Optional[TimerProtocol] = None,
  ) -> None:
      """
      The ``train`` entry point takes in a :class:`~torchtnt.framework.unit.TrainUnit` object, a train dataloader (any Iterable), optional arguments to modify loop execution,
      and runs the training loop.
  
      Args:
          train_unit: an instance of :class:`~torchtnt.framework.unit.TrainUnit` which implements `train_step`.
          train_dataloader: dataloader to be used during training, which can be *any* iterable, including PyTorch DataLoader, DataLoader2, etc.
          max_epochs: the max number of epochs to run. ``None`` means no limit (infinite training) unless stopped by max_steps.
          max_steps: the max number of steps to run. ``None`` means no limit (infinite training) unless stopped by max_epochs.
          max_steps_per_epoch: the max number of steps to run per epoch. None means train until the dataloader is exhausted.
          callbacks: an optional list of :class:`~torchtnt.framework.callback.Callback` s.
          timer: an optional Timer which will be used to time key events (using a Timer with CUDA synchronization may degrade performance).
  
  
      Below is an example of calling :py:func:`~torchtnt.framework.train`.
  
      .. code-block:: python
  
          from torchtnt.framework.train import train
  
          train_unit = MyTrainUnit(module=..., optimizer=..., lr_scheduler=...)
          train_dataloader = torch.utils.data.DataLoader(...)
          train(train_unit, train_dataloader, max_epochs=4)
  
      Below is pseudocode of what the :py:func:`~torchtnt.framework.train` entry point does.
  
      .. code-block:: text
  
          set unit's tracked modules to train mode
          call on_train_start on unit first and then callbacks
          while training is not done:
              while epoch is not done:
                  call on_train_epoch_start on unit first and then callbacks
                  try:
                      call get_next_train_batch on unit
                      call on_train_step_start on callbacks
                      call train_step on unit
                      increment step counter
                      call on_train_step_end on callbacks
                  except StopIteration:
                      break
              increment epoch counter
              call on_train_epoch_end on unit first and then callbacks
          call on_train_end on unit first and then callbacks
      """
      _log_api_usage("train")
      callback_handler = CallbackHandler(callbacks or [])
      state = State(
          entry_point=EntryPoint.TRAIN,
          train_state=PhaseState(
              dataloader=train_dataloader,
              max_epochs=max_epochs,
              max_steps=max_steps,
              max_steps_per_epoch=max_steps_per_epoch,
          ),
          timer=timer,
      )
      try:
          _train_impl(state, train_unit, callback_handler)
          logger.info("Finished train")
          if state.timer:
              logger.info(get_timer_summary(state.timer))
      except Exception as e:
          # TODO: log for diagnostics
          logger.info(f"Exception during train:\n{e}")
          train_unit.on_exception(state, e)
          callback_handler.on_exception(state, train_unit, e)
          raise e
  
file_path: tnt/torchtnt/framework/train.py
imports: import logging
  from typing import Iterable, List, Optional
  import torch
  from pyre_extensions import none_throws
  from torchtnt.framework._callback_handler import CallbackHandler
  from torchtnt.framework._loop_utils import _is_done, _is_epoch_done, _log_api_usage, _maybe_set_distributed_sampler_epoch, _reset_module_training_mode, _set_module_training_mode
  from torchtnt.framework.callback import Callback
  from torchtnt.framework.evaluate import _evaluate_impl
  from torchtnt.framework.state import ActivePhase, EntryPoint, PhaseState, State
  from torchtnt.framework.unit import TTrainData, TTrainUnit
  from torchtnt.framework.utils import get_timing_context
  from torchtnt.utils.timer import get_timer_summary, TimerProtocol
label_desc: Use image node_id 1 for calling a global function with example usage: train(train_unit, train_dataloader) without return types