code: def _predict_impl(
      state: State,
      predict_unit: TPredictUnit,
      callback_handler: CallbackHandler,
  ) -> None:
      # input validation
      predict_state = none_throws(state.predict_state)
  
      state._active_phase = ActivePhase.PREDICT
      logger.info(
          f"Started predict with max_steps_per_epoch={predict_state.max_steps_per_epoch}"
      )
  
      # Set all modules to eval mode
      # access modules made available through AppStateMixin
      tracked_modules = predict_unit.tracked_modules()
      prior_module_train_states = _set_module_training_mode(
          tracked_modules, False
      )
  
      with get_timing_context(
          state, f"{predict_unit.__class__.__name__}.on_predict_start"
      ):
          predict_unit.on_predict_start(state)
      callback_handler.on_predict_start(state, predict_unit)
  
      # Conditionally run this to avoid running this multiple times
      # in the case of resuming from a checkpoint mid-epoch
      if (
          predict_unit.predict_progress.num_steps_completed_in_epoch
          == 0
      ):
          with get_timing_context(
              state,
              f"{predict_unit.__class__.__name__}.on_predict_epoch_start",
          ):
              predict_unit.on_predict_epoch_start(state)
          callback_handler.on_predict_epoch_start(state, predict_unit)
  
      with get_timing_context(state, "predict.iter(dataloader)"):
          data_iter = iter(predict_state.dataloader)
      step_input = data_iter
  
      prev_steps_in_epoch = (
          predict_unit.predict_progress.num_steps_completed_in_epoch
      )
  
      while not (
          state.should_stop
          or _is_epoch_done(
              predict_unit.predict_progress,
              predict_state.max_steps_per_epoch,
              predict_state.max_steps,
          )
      ):
          try:
              with get_timing_context(
                  state, "predict.next(data_iter)"
              ), predict_state.iteration_timer.time("data_wait_time"):
                  step_input = predict_unit.get_next_predict_batch(
                      state, data_iter
                  )
                  callback_handler.on_predict_get_next_batch_end(
                      state, predict_unit
                  )
  
              with predict_state.iteration_timer.time(
                  "predict_iteration_time"
              ):
                  callback_handler.on_predict_step_start(
                      state, predict_unit
                  )
                  predict_state._step_output = (
                      predict_unit.predict_step(state, step_input)
                  )
  
                  predict_unit.predict_progress.increment_step()
                  callback_handler.on_predict_step_end(
                      state, predict_unit
                  )
  
                  # clear step_output to avoid retaining extra memory
                  predict_state._step_output = None
          except StopIteration:
              break
  
      # Possibly warn about an empty dataloader
      any_steps_completed = (
          abs(
              predict_unit.predict_progress.num_steps_completed_in_epoch
              - prev_steps_in_epoch
          )
          > 0
      )
      if not any_steps_completed:
          logger.warning("No steps completed during predict epoch!")
  
      # set progress counters for the next epoch
      predict_unit.predict_progress.increment_epoch()
  
      predict_unit.on_predict_epoch_end(state)
      callback_handler.on_predict_epoch_end(state, predict_unit)
  
      predict_unit.on_predict_end(state)
      callback_handler.on_predict_end(state, predict_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/predict.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_epoch_done, _log_api_usage, _reset_module_training_mode, _set_module_training_mode
  from torchtnt.framework.callback import Callback
  from torchtnt.framework.state import ActivePhase, EntryPoint, PhaseState, State
  from torchtnt.framework.unit import TPredictData, TPredictUnit
  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: _predict_impl(state, predict_unit, callback_handler) without return types