(
self,
low: int,
high: int,
batch_size: int = 1024,
g: torch.Generator = torch.Generator()
)
| 112 | self.tokens = self.tokens[n_tokens:] |
| 113 | |
| 114 | def randint( |
| 115 | self, |
| 116 | low: int, |
| 117 | high: int, |
| 118 | batch_size: int = 1024, |
| 119 | g: torch.Generator = torch.Generator() |
| 120 | ) -> Iterable[int]: |
| 121 | indices = torch.empty(batch_size, dtype=torch.long) |
| 122 | while True: |
| 123 | # record the generator states before sampling |
| 124 | self.rng_state = g.get_state() |
| 125 | indices = torch.randint(low, high, (batch_size,), out=indices, generator=g) |
| 126 | for i in indices[self.rand_id:].tolist(): |
| 127 | self.rand_id += 1 |
| 128 | yield i |
| 129 | self.rand_id = 0 |
| 130 | |
| 131 | def set_epoch(self, epoch): |
| 132 | self._epoch = epoch |
no outgoing calls
no test coverage detected