(self, states)
| 96 | return {'step_counter': tensor.to_numpy(self.step_counter)[0]} |
| 97 | |
| 98 | def set_states(self, states): |
| 99 | self.step_counter = Tensor((1,)) |
| 100 | self.step_counter.set_value(states['step_counter']) |
| 101 | self.lr_value = self.lr(self.step_counter) |
| 102 | |
| 103 | def __call__(self, loss): |
| 104 | self.call(loss) |