code: def test_iteration_time_logger_test_on_train_step_end(self) -> None:
      logger = MagicMock(spec=TensorBoardLogger)
      logger.writer = MagicMock(spec=SummaryWriter)
      state = MagicMock(spec=State)
  
      # Test that the recorded times are tracked separately and that we properly
      # handle when there are no values for a given metric.
      recorded_durations = {
          "train_iteration_time": [1, 3, 5, 7, 9],
          "eval_iteration_time": [],
          "predict_iteration_time": [11, 13, 15, 17, 19],
      }
      state.train_state.iteration_timer.recorded_durations = (
          recorded_durations.copy()
      )
      state.eval_state.iteration_timer.recorded_durations = (
          recorded_durations.copy()
      )
      state.predict_state.iteration_timer.recorded_durations = (
          recorded_durations.copy()
      )
  
      callback = IterationTimeLogger(logger=logger, moving_avg_window=4)
  
      # Set up some dummy units for invoking the callback.
      train_unit = DummyTrainUnit(input_dim=2)
      train_unit.train_progress.increment_step()
      train_unit.train_progress.increment_step()
  
      eval_unit = DummyEvalUnit(input_dim=2)
      eval_unit.eval_progress.increment_step()
      eval_unit.eval_progress.increment_step()
  
      predict_unit = DummyPredictUnit(input_dim=2)
      predict_unit.predict_progress.increment_step()
      predict_unit.predict_progress.increment_step()
  
      # Invoke the callback for each step type we are tracking iteration time for.
      callback = IterationTimeLogger(logger=logger, moving_avg_window=4)
      callback.on_train_step_end(state, train_unit)
      callback.on_eval_step_end(state, eval_unit)
      callback.on_predict_step_end(state, predict_unit)
  
      logger.writer.add_scalar.assert_has_calls(
          [
              call(
                  "Train Iteration Time (seconds)",
                  6.0,  # the average of the last 4 numbers is 6
                  2,  # after incrementing twice, step should be 2
              ),
              call(
                  "Prediction Iteration Time (seconds)",
                  16.0,  # the average of the last 4 numbers is 16
                  2,  # after incrementing twice, step should be 2
              ),
          ]
      )
  
file_path: tnt/tests/framework/callbacks/test_iteration_time_logger.py
imports: import unittest
  from unittest.mock import call, MagicMock
  from torch.utils.tensorboard import SummaryWriter
  from torchtnt.framework._test_utils import DummyEvalUnit, DummyPredictUnit, DummyTrainUnit, generate_random_dataloader
  from torchtnt.framework.callbacks.iteration_time_logger import IterationTimeLogger
  from torchtnt.framework.state import State
  from torchtnt.framework.train import train
  from torchtnt.utils.loggers import TensorBoardLogger
label_desc: Use image node_id 1 for calling the IterationTimeLoggerTest obj's underlying member method code with example usage: obj.test_iteration_time_logger_test_on_train_step_end() without return types