StreamingDataset: Fault Tolerance v2 2/n (#19201)

Co-authored-by: thomas <thomas@thomass-MacBook-Pro.local>

Author: tchaton

src/lightning/data/streaming/dataloader.py

@@ -370,11 +370,10 @@ def __iter__(self) -> Any:
         else:
             yield from super().__iter__()
 
-    def state_dict(self) -> Optional[Dict[str, Any]]:
+    def state_dict(self) -> Dict[str, Any]:
         if isinstance(self.dataset, StreamingDataset):
             assert self.batch_size
-            env = _DistributedEnv.detect()
-            num_samples = self.num_samples_yielded * env.world_size
+            num_samples = self.num_samples_yielded
             return self.dataset.state_dict(num_samples, self.num_workers, self.batch_size)
         return self.dataset.state_dict(self.num_workers, self.batch_size)
 
@@ -384,7 +383,7 @@ def load_state_dict(self, obj: Dict[str, Any]) -> None:
         This is called on each copy of the dataset when resuming.
 
         Args:
-            obj (Dict[str, Any]): The state.
+            obj (Any): The state.
 
         """
         if isinstance(self.dataset, (StreamingDataset, CombinedStreamingDataset)):