Overrides default init by loading value from checkpoint.
(self, *args, **kwargs)
| 28 | |
| 29 | |
| 30 | def _init_from_checkpoint(self, *args, **kwargs): |
| 31 | """Overrides default init by loading value from checkpoint.""" |
| 32 | # pylint: disable=protected-access |
| 33 | self._old_init(*args, **kwargs) |
| 34 | ckpt_name = self._map_func(self._shared_name) |
| 35 | if ckpt_name not in self._ckpt_var_cache: |
| 36 | raise errors.NotFoundError(None, None, |
| 37 | "%s not found in checkpoint" % ckpt_name) |
| 38 | |
| 39 | val = self._ckpt_var_cache.get(ckpt_name, None) |
| 40 | if val is not None: |
| 41 | self.assign(val) |
| 42 | # Avoid assigning for the second time. |
| 43 | self._ckpt_var_cache[ckpt_name] = None |
| 44 | # pylint: enable=protected-access |
| 45 | |
| 46 | |
| 47 | @contextlib.contextmanager |