code: def test_initialization_and_get_next_batch(self) -> None:
      unit = TestUnit()
      self.assertIsNotNone(unit.train_progress)
      self.assertIsNotNone(unit.eval_progress)
      self.assertIsNotNone(unit.predict_progress)
  
      tensor_1 = torch.ones(1)
      tensor_2 = torch.zeros(1)
      state = get_dummy_train_state()
  
      # test train next batch - exepct to return the elements within the iterable
      train_data_iter = iter([tensor_1, tensor_2])
      self.assertEqual(
          unit.get_next_train_batch(state, train_data_iter), tensor_1
      )
      self.assertEqual(
          unit.get_next_train_batch(state, train_data_iter), tensor_2
      )
  
      # test predict next batch - exepct to return the elements within the iterable
      self.assertEqual(
          unit.get_next_predict_batch(
              state, iter([tensor_1, tensor_2])
          ),
          tensor_1,
      )
  
      # test eval next batch - exepct to return the iterable
      data_iter = iter([tensor_1, tensor_2])
      next_eval_batch = unit.get_next_eval_batch(state, data_iter)
      self.assertEqual(next_eval_batch, data_iter)
  
file_path: tnt/tests/framework/test_unit.py
imports: import unittest
  from typing import Iterator
  import torch
  from torchtnt.framework._test_utils import get_dummy_train_state
  from torchtnt.framework.state import State
  from torchtnt.framework.unit import EvalUnit, PredictUnit, TrainUnit
label_desc: Use image node_id 1 for calling the UnitTest obj's underlying member method code with example usage: obj.test_initialization_and_get_next_batch() without return types