code: def _train_impl(
      state: State,
      train_unit: TTrainUnit,
      callback_handler: CallbackHandler,
  ) -> None:
      train_state = none_throws(state.train_state)
  
      logger.info(
          f"Started train with max_epochs={train_state.max_epochs}, max_steps={train_state.max_steps}, max_steps_per_epoch={train_state.max_steps_per_epoch}"
      )
      if (
          train_state.max_epochs is None
          and train_state.max_steps is None
      ):
          logger.warning(
              "Will run infinite training, since both max_epochs and max_steps were not set."
          )
      state._active_phase = ActivePhase.TRAIN
  
      # Set all modules to train() mode
      # access modules made available through AppStateMixin
      tracked_modules = train_unit.tracked_modules()
      prior_module_train_states = _set_module_training_mode(
          tracked_modules, True
      )
  
      train_unit.on_train_start(state)
      callback_handler.on_train_start(state, train_unit)
  
      while not (
          state.should_stop
          or _is_done(
              train_unit.train_progress,
              train_state.max_epochs,
              train_state.max_steps,
          )
      ):
          _train_epoch_impl(state, train_unit, callback_handler)
          logger.info(
              "After train epoch, train progress: "
              f"num_epochs_completed = {train_unit.train_progress.num_epochs_completed}, "
              f"num_steps_completed = {train_unit.train_progress.num_steps_completed}"
          )
  
      train_unit.on_train_end(state)
      callback_handler.on_train_end(state, train_unit)
  
      # Reset training mode for modules at the end of the epoch
      # This ensures that side-effects made by the loop are reset before
      # returning back to the user
      _reset_module_training_mode(
          tracked_modules, prior_module_train_states
      )
  
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 2 for calling a global function with example usage: _train_impl(state, train_unit, callback_handler) without return types