(self, states)
| 207 | return states |
| 208 | |
| 209 | def set_states(self, states): |
| 210 | super().set_states(states) |
| 211 | if 'moments' in states: |
| 212 | self.moments = states['moments'] |
| 213 | self.mom_value = self.momentum(self.step_counter) |
| 214 | |
| 215 | |
| 216 | # Data augmentation |