| 20 | |
| 21 | |
| 22 | class State(object): |
| 23 | def __init__(self): |
| 24 | self.epoch = 1 |
| 25 | self.iteration = 0 |
| 26 | self.dataloader = None |
| 27 | self.model = None |
| 28 | self.optimizer = None |
| 29 | |
| 30 | def register(self, **kwargs): |
| 31 | for k, v in kwargs.items(): |
| 32 | assert k in ["epoch", "iteration", "dataloader", "model", "optimizer"] |
| 33 | setattr(self, k, v) |
| 34 | |
| 35 | |
| 36 | class Engine(object): |