code: def test_op_saves_and_restores_its_state(self) -> None:
      pipeline = DataPipeline.count(start=4).take(10).and_return()
  
      d = None
  
      it = iter(pipeline)
  
      # Move to the fifth example.
      for _ in range(5):
          d = next(it)
  
      assert d == 8
  
      state_dict = pipeline.state_dict()
  
      # Read a few examples before we roll back.
      for _ in range(4):
          d = next(it)
  
      assert d == 12
  
      # Expected to roll back to the fifth example.
      pipeline.load_state_dict(state_dict)
  
      # Move to EOD.
      for _ in range(5):
          d = next(it)
  
      assert d == 13
  
      state_dict = pipeline.state_dict()
  
      pipeline.reset()
  
      # Expected to be EOD.
      pipeline.load_state_dict(state_dict)
  
      with pytest.raises(StopIteration):
          next(iter(pipeline))
  
imports: import pytest
  from fairseq2.data import DataPipeline
label_desc: Use image node_id 2 for calling the TestCountOp obj's underlying member method code with example usage: obj.test_op_saves_and_restores_its_state() without return types