(self)
| 134 | self.dataset.set_epoch(epoch) |
| 135 | |
| 136 | def state_dict(self): |
| 137 | return { |
| 138 | 'states': self.states, |
| 139 | 'buffer': self.buffer.clone(), |
| 140 | 'tokens': deepcopy(self.tokens), |
| 141 | 'rand_id': self.rand_id, |
| 142 | 'token_id': self.token_id, |
| 143 | 'rng_state': self.rng_state, |
| 144 | 'epoch': self._epoch |
| 145 | } |
| 146 | |
| 147 | def load_state_dict(self, state_dict): |
| 148 | self.states = state_dict['states'] |