Bugfix: inconsistent streaming dataloader state (specific to StreamingDataset)

src/litdata/streaming/dataloader.py

@@ -615,6 +615,7 @@ def __iter__(self) -> Any:
             self.current_epoch += 1
             self._num_samples_yielded_combined = {}
             self._num_samples_yielded_streaming = 0
+            self.dataset.reset_state_dict()
 
         self.dataset.set_epoch(self.current_epoch)
 

src/litdata/streaming/dataset.py

@@ -388,7 +388,7 @@ def state_dict(self, num_samples_yielded: int, num_workers: int, batch_size: int
 
         return {
             "num_samples_yielded": num_samples_yielded,
-            "num_workers": num_workers,
+            "num_workers": num_workers if num_workers > 0 else 1,
             "batch_size": batch_size,
             "current_epoch": self.current_epoch,
             "input_dir_path": self.input_dir.path,
@@ -407,6 +407,9 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             # the state is restored within the workers
             self._state_dict = state_dict
 
+    def reset_state_dict(self) -> None:
+        self._state_dict = None
+
     def _validate_state_dict(self) -> None:
         assert self._state_dict
         assert self.worker_env