code: def test_combined_dataset_state_dict():
      dataset = CombinedStreamingDataset(
          [TestStatefulDataset(10, 1), TestStatefulDataset(10, -1)],
          42,
          weights=(0.5, 0.5),
      )
      assert dataset.state_dict(0, 1) == {}
      dataset_iter = iter(dataset)
      assert dataset.state_dict(0, 1) == {
          "0": {"counter": 0},
          "1": {"counter": 0},
      }
  
      dataset2 = CombinedStreamingDataset(
          [TestStatefulDataset(10, 1), TestStatefulDataset(10, -1)],
          42,
          weights=(0.5, 0.5),
      )
      assert dataset2.state_dict(0, 1) == {}
  
      data = []
      states = []
      for i, value in enumerate(dataset_iter):
          state = dataset.state_dict(i, 1)
          data.append(value)
          states.append(state)
  
      assert data == [
          0,
          0,
          1,
          2,
          -1,
          -2,
          -3,
          3,
          4,
          5,
          6,
          -4,
          7,
          8,
          -5,
          -6,
          9,
          -7,
          -8,
      ]
      assert states == [
          {"0": {"counter": 0}, "1": {"counter": 1}},
          {"0": {"counter": 1}, "1": {"counter": 1}},
          {"0": {"counter": 2}, "1": {"counter": 1}},
          {"0": {"counter": 3}, "1": {"counter": 1}},
          {"0": {"counter": 3}, "1": {"counter": 2}},
          {"0": {"counter": 3}, "1": {"counter": 3}},
          {"0": {"counter": 3}, "1": {"counter": 4}},
          {"0": {"counter": 4}, "1": {"counter": 4}},
          {"0": {"counter": 5}, "1": {"counter": 4}},
          {"0": {"counter": 6}, "1": {"counter": 4}},
          {"0": {"counter": 7}, "1": {"counter": 4}},
          {"0": {"counter": 7}, "1": {"counter": 5}},
          {"0": {"counter": 8}, "1": {"counter": 5}},
          {"0": {"counter": 9}, "1": {"counter": 5}},
          {"0": {"counter": 9}, "1": {"counter": 6}},
          {"0": {"counter": 9}, "1": {"counter": 7}},
          {"0": {"counter": 10}, "1": {"counter": 7}},
          {"0": {"counter": 10}, "1": {"counter": 8}},
          {"0": {"counter": 10}, "1": {"counter": 9}},
      ]
  
      dataset2 = CombinedStreamingDataset(
          [TestStatefulDataset(10, 1), TestStatefulDataset(10, -1)],
          42,
          weights=(0.5, 0.5),
      )
      assert dataset2.state_dict(0, 1) == {}
      dataset2_iter = iter(dataset2)
  
      data_2 = []
      for state in states:
          dataset.load_state_dict(state)
          data_2.append(next(dataset2_iter))
  
      assert data == data_2
  
file_path: pytorch-lightning/tests/tests_data/streaming/test_combined.py
imports: import pytest
  from lightning.data.streaming.combined import CombinedStreamingDataset
  from torch.utils.data import IterableDataset
label_desc: Use image node_id 2 for calling a global function with example usage: test_combined_dataset_state_dict() without return types