docstr: Test fit entry point with a callback
docstr_tok: ["Test","fit","entry","point","with","a","callback"]
code: def test_fit_with_callback(self) -> None:
      """
      Test fit entry point with a callback
      """
      input_dim = 2
      train_dataset_len = 10
      eval_dataset_len = 6
      batch_size = 2
      max_epochs = 4
      expected_num_total_train_steps = (
          train_dataset_len / batch_size * max_epochs
      )
      expected_num_total_eval_steps = (
          eval_dataset_len / batch_size * max_epochs
      )
  
      my_unit = DummyFitUnit(2)
      train_dataloader = generate_random_dataloader(
          train_dataset_len, input_dim, batch_size
      )
      eval_dataloader = generate_random_dataloader(
          eval_dataset_len, input_dim, batch_size
      )
  
      callback_mock = MagicMock(spec=Callback)
  
      fit(
          my_unit,
          train_dataloader=train_dataloader,
          eval_dataloader=eval_dataloader,
          max_epochs=max_epochs,
          callbacks=[callback_mock],
      )
  
      self.assertEqual(callback_mock.on_train_start.call_count, 1)
      self.assertEqual(
          callback_mock.on_train_epoch_start.call_count, max_epochs
      )
      self.assertEqual(
          callback_mock.on_train_step_start.call_count,
          expected_num_total_train_steps,
      )
      self.assertEqual(
          callback_mock.on_train_step_end.call_count,
          expected_num_total_train_steps,
      )
      self.assertEqual(
          callback_mock.on_train_epoch_end.call_count, max_epochs
      )
      self.assertEqual(callback_mock.on_train_end.call_count, 1)
  
      self.assertEqual(
          callback_mock.on_eval_start.call_count, max_epochs
      )
      self.assertEqual(
          callback_mock.on_eval_epoch_start.call_count, max_epochs
      )
      self.assertEqual(
          callback_mock.on_eval_step_start.call_count,
          expected_num_total_eval_steps,
      )
      self.assertEqual(
          callback_mock.on_eval_step_end.call_count,
          expected_num_total_eval_steps,
      )
      self.assertEqual(
          callback_mock.on_eval_epoch_end.call_count, max_epochs
      )
      self.assertEqual(callback_mock.on_eval_end.call_count, max_epochs)
  
imports: import math
  import unittest
  from typing import Tuple
  from unittest.mock import MagicMock
  import torch
  from torch import nn
  from torchtnt.framework._test_utils import DummyFitUnit, generate_random_dataloader
  from torchtnt.framework.callback import Callback
  from torchtnt.framework.fit import fit
  from torchtnt.framework.state import ActivePhase, State
  from torchtnt.framework.unit import EvalUnit, TrainUnit, TTrainUnit
  from torchtnt.utils.timer import Timer
label_desc: Use image node_id 5 for calling the FitTest obj's underlying member method code with example usage: obj.test_fit_with_callback() without return types