code: def _train_epoch_impl(
      state: State,
      train_unit: TTrainUnit,
      callback_handler: CallbackHandler,
  ) -> None:
      logger.info("Started train epoch")
      state._active_phase = ActivePhase.TRAIN
  
      train_state = none_throws(state.train_state)
  
      evaluate_every_n_steps = None
      evaluate_every_n_epochs = None
      if state.eval_state:
          if state.eval_state.evaluate_every_n_steps:
              evaluate_every_n_steps = (
                  state.eval_state.evaluate_every_n_steps
              )
          if state.eval_state.evaluate_every_n_epochs:
              evaluate_every_n_epochs = (
                  state.eval_state.evaluate_every_n_epochs
              )
  
      # Check the progress to conditionally run this
      # to avoid running this multiple times
      # in the case of resuming from a checkpoint mid-epoch
      if train_unit.train_progress.num_steps_completed_in_epoch == 0:
          train_unit.on_train_epoch_start(state)
          callback_handler.on_train_epoch_start(state, train_unit)
  
      _maybe_set_distributed_sampler_epoch(
          train_state.dataloader,
          train_unit.train_progress.num_epochs_completed,
      )
  
      with get_timing_context(state, "train.iter(dataloader)"):
          data_iter = iter(train_state.dataloader)
  
      prev_steps_in_epoch = (
          train_unit.train_progress.num_steps_completed_in_epoch
      )
  
      while not (
          state.should_stop
          or _is_epoch_done(
              train_unit.train_progress,
              train_state.max_steps_per_epoch,
              train_state.max_steps,
          )
      ):
          try:
              with get_timing_context(
                  state, "train.next(data_iter)"
              ), train_state.iteration_timer.time("data_wait_time"):
                  step_input = train_unit.get_next_train_batch(
                      state, data_iter
                  )
                  callback_handler.on_train_get_next_batch_end(
                      state, train_unit
                  )
  
              with train_state.iteration_timer.time(
                  "train_iteration_time"
              ):
                  callback_handler.on_train_step_start(
                      state, train_unit
                  )
                  train_state._step_output = train_unit.train_step(
                      state, step_input
                  )
                  train_unit.train_progress.increment_step()
                  callback_handler.on_train_step_end(state, train_unit)
  
                  # clear step_output to avoid retaining extra memory
                  train_state._step_output = None
  
              if (
                  evaluate_every_n_steps
                  and train_unit.train_progress.num_steps_completed
                  % evaluate_every_n_steps
                  == 0
              ):
                  _evaluate_impl(
                      state,
                      # pyre-ignore: Incompatible parameter type [6]
                      train_unit,
                      callback_handler,
                  )
                  logger.info(
                      "Finished evaluation. Resuming training epoch"
                  )
                  state._active_phase = ActivePhase.TRAIN
  
          except StopIteration:
              logger.info("Reached end of train dataloader")
              break
  
      # Possibly warn about an empty dataloader
      any_steps_completed = (
          abs(
              train_unit.train_progress.num_steps_completed_in_epoch
              - prev_steps_in_epoch
          )
          > 0
      )
      if not any_steps_completed:
          logger.warning("No steps completed during train epoch!")
  
      # set progress counters for the next epoch
      train_unit.train_progress.increment_epoch()
  
      train_unit.on_train_epoch_end(state)
      callback_handler.on_train_epoch_end(state, train_unit)
  
      if (
          evaluate_every_n_epochs
          and train_unit.train_progress.num_epochs_completed
          % evaluate_every_n_epochs
          == 0
      ):
          _evaluate_impl(
              state,
              # pyre-ignore: Incompatible parameter type [6]
              train_unit,
              callback_handler,
          )
          state._active_phase = ActivePhase.TRAIN
  
      logger.info("Ended train epoch")
  
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 3 for calling a global function with example usage: _train_epoch_impl(state, train_unit, callback_handler) without return types