| 76 | """ |
| 77 | |
| 78 | def __init__(self, lr, dtype=tensor.float32): |
| 79 | # init lr(could be a constant scalar or a learning rate scheduler) |
| 80 | if type(lr) == float or type(lr) == int: |
| 81 | self.lr = Constant(lr) |
| 82 | elif isinstance(lr, DecayScheduler): |
| 83 | self.lr = lr |
| 84 | else: |
| 85 | raise TypeError("Wrong learning rate type") |
| 86 | |
| 87 | # init step counter |
| 88 | self.dtype = dtype |
| 89 | # TODO change type to int32 |
| 90 | self.step_counter = Tensor((1,), dtype=tensor.float32) |
| 91 | self.step_counter.set_value(0) |
| 92 | self.lr_value = self.lr(self.step_counter) |
| 93 | |
| 94 | def get_states(self): |
| 95 | # skip DecayScheduler as it does not have persistent states |