| 20 | |
| 21 | |
| 22 | def monkey_patch_program(): |
| 23 | @signature_safe_contextmanager |
| 24 | def _lr_schedule_guard(self, is_with_opt=False): |
| 25 | # TODO(dev): Currently there has not equivalent of op_role in PIR |
| 26 | # mode, so we simply remove _lr_schedule_guard here, this should |
| 27 | # be fixed in the future. |
| 28 | yield |
| 29 | |
| 30 | def state_dict(self, mode="all", scope=None): |
| 31 | from paddle.base import core |
| 32 | from paddle.base.executor import global_scope |
| 33 | |
| 34 | if scope is not None and not isinstance(scope, core._Scope): |
| 35 | raise TypeError( |
| 36 | f"`scope` should be None or `paddle.static.Scope'` type, but received {type(scope)}." |
| 37 | ) |
| 38 | |
| 39 | if scope is None: |
| 40 | scope = global_scope() |
| 41 | return self._state_dict(mode, scope) |
| 42 | |
| 43 | program_attrs = { |
| 44 | "_lr_schedule_guard": _lr_schedule_guard, |
| 45 | "state_dict": state_dict, |
| 46 | } |
| 47 | |
| 48 | global _already_patch_program |
| 49 | if not _already_patch_program: |
| 50 | for attr, value in program_attrs.items(): |
| 51 | setattr(Program, attr, value) |
| 52 | |
| 53 | _already_patch_program = True |