update generator usage

pytorch · Feb 11, 2025 · 0a90c04 · 0a90c04
1 parent 10ba260
commit 0a90c04
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/torchdata/stateful_dataloader/sampler.py b/torchdata/stateful_dataloader/sampler.py
@@ -15,7 +15,7 @@
 from .stateful import Stateful
 
 
-class StatefulRandomSamplerIterator(Iterator[list[int]], Stateful):
+class StatefulRandomSamplerIterator(Iterator[int], Stateful):
     _GENERATOR = "generator"
     _YIELDED = "yielded"
 
@@ -25,7 +25,6 @@ def __init__(self, sampler):
         self.yielded = 0
         self.next_yielded = None
         self.n = len(sampler.data_source)
-        self.generator = sampler.generator
         self.replacement = sampler.replacement
         self.num_samples = sampler.num_samples
         self.chunk_size = 32
@@ -46,7 +45,7 @@ def __next__(self):
                         high=self.n,
                         size=(self.chunk_size,),
                         dtype=torch.int64,
-                        generator=self.generator,
+                        generator=self.sampler.generator,
                     ).tolist()
                     self.perm_index = 0
                 value = self.perm[self.perm_index]
@@ -62,7 +61,7 @@ def __next__(self):
                         high=self.n,
                         size=(remainder,),
                         dtype=torch.int64,
-                        generator=self.generator,
+                        generator=self.sampler.generator,
                     ).tolist()
                     self.perm_index = 0
                 value = self.perm[self.perm_index]
@@ -78,7 +77,7 @@ def __next__(self):
             remainder = self.num_samples % self.n
             if self.chunk_index < num_full_perms:
                 if self.perm is None or not self.perm:
-                    self.perm = torch.randperm(self.n, generator=self.generator).tolist()
+                    self.perm = torch.randperm(self.n, generator=self.sampler.generator).tolist()
                     self.perm_index = 0
                 value = self.perm[self.perm_index]
                 self.perm_index += 1
@@ -89,7 +88,7 @@ def __next__(self):
                 return value
             elif remainder > 0:
                 if self.perm is None or not self.perm:
-                    self.perm = torch.randperm(self.n, generator=self.generator).tolist()[:remainder]
+                    self.perm = torch.randperm(self.n, generator=self.sampler.generator).tolist()[:remainder]
                     self.perm_index = 0
                 value = self.perm[self.perm_index]
                 self.perm_index += 1
@@ -109,7 +108,7 @@ def state_dict(self) -> dict:
     def load_state_dict(self, state_dict: dict) -> None:
         self.next_yielded = state_dict[self._YIELDED]
         self.generator_state = state_dict[self._GENERATOR]
-        self.generator.set_state(self.generator_state)
+        self.sampler.generator.set_state(self.generator_state)
         if self.next_yielded is not None:
             for _ in range(self.next_yielded - self.yielded):
                 next(self)