docstr: Test IterationTimeLogger callback with train entry point
code: def test_with_train_epoch(self) -> None:
      """
      Test IterationTimeLogger callback with train entry point
      """
  
      my_unit = DummyTrainUnit(input_dim=2)
      logger = MagicMock(spec=TensorBoardLogger)
      logger.writer = MagicMock(spec=SummaryWriter)
      callback = IterationTimeLogger(
          logger, moving_avg_window=1, log_every_n_steps=3
      )
      dataloader = generate_random_dataloader(
          num_samples=12, input_dim=2, batch_size=2
      )
      train(my_unit, dataloader, max_epochs=2, callbacks=[callback])
      # 2 epochs, 6 iterations each, logging every third step
      self.assertEqual(logger.writer.add_scalar.call_count, 4)
  
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 2 for calling the IterationTimeLoggerTest obj's underlying member method code with example usage: obj.test_with_train_epoch() without return types